d9d 0.1.0__py3-none-any.whl

d9d/loop/run/train.py

@@ -0,0 +1,355 @@
+from pathlib import Path
+
+from tqdm import tqdm
+
+from d9d.core.dist_context import DeviceMeshParameters
+from d9d.internals.determinism import set_seeds
+from d9d.internals.pipeline_state import PipelineStateHandler
+from d9d.loop.component import (
+    BatchMaths,
+    DataLoaderFactory,
+    GradientClipper,
+    GradientManager,
+    JobLogger,
+    JobProfiler,
+    LossComputer,
+    ManualGarbageCollector,
+    ModelStageExporter,
+    ModelStageFactory,
+    OptimizerFactory,
+    StateCheckpointer,
+    Stepper,
+    TimeoutManager,
+    TrainTaskOperator,
+)
+from d9d.loop.config import TrainerConfig
+from d9d.loop.control import (
+    CreateMetricsContext,
+    DatasetProvider,
+    FinalizeContext,
+    LRSchedulerProvider,
+    ModelProvider,
+    OptimizerProvider,
+    TrainTaskProvider,
+    TrainTaskProviderContext,
+)
+from d9d.loop.state import TrainJobState
+from d9d.metric.impl import ComposeMetric
+
+
+class TrainingConfigurator:
+    """
+    Orchestrates the assembly of the distributed training environment.
+
+    This class binds the infrastructure configuration (DeviceMesh), the training
+    parameters (TrainerConfig), and the user-defined logic (Providers) to create
+    a fully initialized state object capable of running the training loop.
+    """
+
+    def __init__(
+            self,
+            mesh: DeviceMeshParameters,
+            parameters: TrainerConfig,
+            task_provider: TrainTaskProvider,
+            model_provider: ModelProvider,
+            data_provider: DatasetProvider,
+            optimizer_provider: OptimizerProvider,
+            lr_scheduler_provider: LRSchedulerProvider
+    ):
+        """
+        Constructs a configurator capable of building the full training state.
+
+        Args:
+            mesh: Definition of the distributed device mesh topology.
+            parameters: The global configuration object for the trainer.
+            task_provider: Factory for creating the training task logic.
+            model_provider: Factory for defining and creating model stages.
+            data_provider: Factory for providing training datasets.
+            optimizer_provider: Factory for creating the optimizer.
+            lr_scheduler_provider: Factory for creating the learning rate scheduler.
+        """
+        self._mesh = mesh
+        self._parameters = parameters
+        self._task_provider = task_provider
+        self._model_provider = model_provider
+        self._data_provider = data_provider
+        self._optimizer_provider = optimizer_provider
+        self._lr_scheduler_provider = lr_scheduler_provider
+
+    def _build_new_training_state(self) -> TrainJobState:
+        dist_context = self._mesh.build()
+
+        set_seeds(dist_context, seed=self._parameters.determinism.base_seed)
+
+        task = self._task_provider(TrainTaskProviderContext(
+            dist_context=dist_context
+        ))
+
+        batch_maths = BatchMaths(
+            dist_context=dist_context,
+            config_batching=self._parameters.batching,
+            config_pipelining=self._parameters.pipelining
+        )
+
+        data_loader_factory = DataLoaderFactory(
+            dist_context=dist_context,
+            provider=self._data_provider,
+            config_data_loading=self._parameters.data_loading,
+            batch_maths=batch_maths
+        )
+        data_loader_train = data_loader_factory.build_dataloader_for_train_job()
+
+        stepper = Stepper(
+            initial_step=1,
+            total_steps=len(data_loader_train)
+        )
+
+        pipeline_state_handler = PipelineStateHandler(
+            sharding_spec={},
+            num_shards=batch_maths.num_microbatches_pipelining
+        )
+
+        loss_computer = LossComputer(
+            state=pipeline_state_handler,
+            task=task,
+            stepper=stepper
+        )
+
+        schedule, modules = ModelStageFactory(
+            model_provider=self._model_provider,
+            dist_context=dist_context,
+            config_model=self._parameters.model_stage_factory,
+            config_pipelining=self._parameters.pipelining,
+            batch_maths=batch_maths,
+            loss_computer=loss_computer
+        ).build_pipeline_and_modules()
+
+        metrics = ComposeMetric(task.create_metrics(CreateMetricsContext()).metrics)
+        metrics.to("cuda")
+
+        task_operator = TrainTaskOperator(
+            dist_context=dist_context,
+            task=task,
+            pp_schedule=schedule,
+            tracked_modules=modules,
+            loss_computer=loss_computer,
+            pipeline_state=pipeline_state_handler,
+            metrics=metrics
+        )
+
+        grad_clipper = GradientClipper(
+            dist_context=dist_context,
+            tracked_modules=modules,
+            config=self._parameters.gradient_clipping,
+            stepper=stepper
+        )
+
+        optimizer, scheduler = OptimizerFactory(
+            dist_context=dist_context,
+            tracked_modules=modules,
+            optimizer_provider=self._optimizer_provider,
+            stepper=stepper,
+            lr_scheduler_provider=self._lr_scheduler_provider
+        ).build_optimizer_and_scheduler()
+
+        gc = ManualGarbageCollector(
+            dist_ctx=dist_context,
+            config=self._parameters.gc,
+            step=stepper
+        )
+
+        checkpointer = StateCheckpointer(
+            dist_context=dist_context,
+            stepper=stepper,
+            config=self._parameters.checkpointing,
+            gc=gc,
+            run_name=self._parameters.run.name
+        )
+
+        profiler = JobProfiler(
+            dist_context=dist_context,
+            stepper=stepper,
+            config=self._parameters.profiling
+        )
+
+        exporter = ModelStageExporter(
+            model_provider=self._model_provider,
+            dist_context=dist_context,
+            modules=modules
+        )
+
+        gradient_manager = GradientManager(
+            dist_context=dist_context,
+            tracked_modules=modules,
+            batch_maths=batch_maths,
+            config=self._parameters.gradient_manager
+        )
+
+        timeout_manager = TimeoutManager(
+            dist_context=dist_context,
+            config=self._parameters.timeout
+        )
+
+        job_logger = JobLogger(
+            dist_context=dist_context,
+            config=self._parameters.logging,
+            metrics=metrics,
+            stepper=stepper,
+            run_config=self._parameters.run,
+            additional_hparams={
+                "task": task.dump_hparams(),
+                "model": self._model_provider.dump_hparams()
+            }
+        )
+
+        return TrainJobState(
+            dist_context=dist_context,
+            data_loader=data_loader_train,
+            stepper=stepper,
+            tracked_modules=modules,
+            garbage_collector=gc,
+            batch_maths=batch_maths,
+            checkpointer=checkpointer,
+            optimizer=optimizer,
+            task=task,
+            lr_scheduler=scheduler,
+            gradient_clipper=grad_clipper,
+            profiler=profiler,
+            exporter=exporter,
+            metrics=metrics,
+            logger=job_logger,
+            gradient_manager=gradient_manager,
+            timeout_manager=timeout_manager,
+            task_operator=task_operator
+        )
+
+    def configure(self) -> "Trainer":
+        """
+        Instantiates all training components and returns a configured Trainer.
+
+        This method triggers the creation of the distributed context, sets seeds,
+        builds the model, optimizer, data loaders, and attaches all auxiliary
+        components (logging, profiling, checkpointing).
+
+        Returns:
+            Trainer: A ready-to-use trainer instance encapsulating the job state.
+        """
+        state = self._build_new_training_state()
+
+        return Trainer(state)
+
+
+class Trainer:
+    """
+    The main execution engine for running a distributed training job.
+
+    This class manages the training loop, lifecycle events, distributed synchronization,
+    and periodic side-effects (logging, checkpointing).
+    """
+
+    def __init__(self, state: TrainJobState):
+        """
+        Constructs a Trainer from a pre-built job state.
+
+        Args:
+            state: The encapsulated state object containing all initialized
+                components (model, optimizer, dist_context, etc.).
+        """
+        self._state = state
+
+    def train(self):
+        """
+        Executes the full training workflow.
+        """
+        self._state.dist_context.logger.info("Waiting for the world to start training")
+        self._state.dist_context.wait_world()
+        self._state.dist_context.logger.info("Trying to load last checkpoint before doing anything else")
+        self._state.checkpointer.load_last_checkpoint(self._state)
+
+        if self._state.stepper.current_step >= self._state.stepper.total_steps:
+            self._state.dist_context.logger.info("Already trained fully, will do nothing")
+            return
+
+        self._state.dist_context.wait_world()
+
+        with (
+            tqdm(
+                desc="Training",
+                total=self._state.stepper.total_steps,
+                disable=not self._state.dist_context.is_local_main_process,
+                initial=self._state.stepper.current_step
+            ) as bar,
+            self._state.logger.new_run() as run,
+            self._state.garbage_collector as gc,
+            self._state.profiler.open() as profiler,
+            self._state.gradient_manager.install(),
+            self._state.gradient_clipper.install()
+        ):
+            self._state.timeout_manager.step()
+            run.set_context({"stage": "train"})
+
+            for batch_group in self._state.data_loader:
+                run.set_step(self._state.stepper.current_step)
+
+                for batch in batch_group:
+                    # we do both forward and backward passes
+                    # since GradientManager is installed - it should start performing
+                    # synchronization overlapping grad sync with compute
+                    loss = self._state.task_operator.forward_backward(batch)
+
+                    # add loss for grad manager - it want it for grad reduction
+                    if loss is not None:
+                        self._state.gradient_manager.add_loss_with_weight(loss.loss, loss.loss_weight)
+
+                # metrics were successfully accumulated during forward passes - we can schedule their synchronization
+                self._state.logger.trigger_sync()
+
+                # wait for gradient synchronization finishes and scale them
+                self._state.gradient_manager.sync_and_scale()
+
+                # clip grads after they are synced across world
+                self._state.gradient_clipper.clip_and_log(run)
+
+                # optimize (it won't sync grads - they are already Replicate-d)
+                self._state.optimizer.step()
+
+                # update LR
+                self._state.lr_scheduler.step()
+
+                # log everything
+                self._state.logger.log(
+                    run,
+                    loss_value=self._state.gradient_manager.compute_global_loss()
+                )
+
+                # reset grads
+                self._state.gradient_manager.zero_grad()
+
+                gc.collect_periodic()
+                self._state.stepper.step()
+                bar.update()
+
+                # checkpoint at the end of the step
+                self._state.checkpointer.checkpoint_if_needed(self._state)
+
+                if profiler:
+                    profiler.step()
+
+            self._state.task.finalize(FinalizeContext())
+
+    def export(self, export_to: Path, load_checkpoint: bool):
+        """
+        Exports the current model state to the specified directory.
+
+        This handles the distributed saving logic, allowing the model to be
+        reconstituted later or used for inference.
+
+        Args:
+            export_to: The directory path where the model artifacts will be saved.
+            load_checkpoint: If True, attempts to load the latest checkpoint
+                into the model before exporting.
+        """
+        if load_checkpoint:
+            self._state.checkpointer.load_last_checkpoint(self._state)
+
+        self._state.exporter.export(export_to)

d9d/loop/state.py

@@ -0,0 +1,143 @@
+import dataclasses
+from typing import Any
+
+from torch.distributed.checkpoint.stateful import Stateful
+from torchdata.stateful_dataloader import StatefulDataLoader
+
+from d9d.core.dist_context import DistributedContext
+from d9d.core.protocol import LRSchedulerProtocol, OptimizerProtocol
+from d9d.loop.component import (
+    BatchMaths,
+    GradientClipper,
+    GradientManager,
+    JobLogger,
+    JobProfiler,
+    ManualGarbageCollector,
+    ModelStageExporter,
+    StateCheckpointer,
+    Stepper,
+    TimeoutManager,
+    TrackedModules,
+    TrainTaskOperator,
+)
+from d9d.loop.control import InferenceTask, TrainTask
+from d9d.metric.impl import ComposeMetric
+
+
+@dataclasses.dataclass(kw_only=True)
+class JobState(Stateful):
+    """
+    Base container for the state of a distributed execution job.
+
+    This dataclass holds the common infrastructure components required for both
+    training and inference loops. It implements the Stateful protocol to support
+    checkpointing of its internal components.
+
+    Attributes:
+        dist_context: The distributed context.
+        stepper: Component for tracking the current global step and total steps.
+        garbage_collector: Component for manual control of Python garbage collection.
+        checkpointer: Component responsible for saving and loading execution states.
+        profiler: Component for performance profiling.
+        tracked_modules: Container holding the model (or model parts) being executed.
+        batch_maths: Helper for calculating batch sizes and gradient accumulation steps.
+        data_loader: The input data stream.
+        timeout_manager: Component for checking and refreshing distributed timeouts.
+    """
+    dist_context: DistributedContext
+
+    stepper: Stepper
+    garbage_collector: ManualGarbageCollector
+    checkpointer: StateCheckpointer
+    profiler: JobProfiler
+
+    tracked_modules: TrackedModules
+    batch_maths: BatchMaths
+
+    data_loader: StatefulDataLoader
+
+    timeout_manager: TimeoutManager
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "stepper": self.stepper.state_dict(),
+            "tracked_modules": self.tracked_modules.state_dict(),
+            "data_loader": self.data_loader.state_dict()
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        self.stepper.load_state_dict(state_dict["stepper"])
+        self.tracked_modules.load_state_dict(state_dict["tracked_modules"])
+        self.data_loader.load_state_dict(state_dict["data_loader"])
+
+
+@dataclasses.dataclass(kw_only=True)
+class TrainJobState(JobState):
+    """
+    Container for the state of a training job.
+
+    Extends JobState to include components specific to training, such as
+    optimization, gradient management, and loss computation.
+
+    Attributes:
+        task: The specific training task logic definition.
+        gradient_manager: Component handling gradient synchronization.
+        metrics: Container for aggregating training metrics.
+        task_operator: Executor for running forward and backward passes.
+        logger: Component for logging metrics and system status.
+        optimizer: The optimizer instance updating model parameters.
+        lr_scheduler: The scheduler adjusting the learning rate.
+        gradient_clipper: Component for clipping gradient norms.
+        exporter: Component for exporting the final model artifacts.
+    """
+    task: TrainTask
+    gradient_manager: GradientManager
+    metrics: ComposeMetric
+    task_operator: TrainTaskOperator
+
+    logger: JobLogger
+
+    optimizer: OptimizerProtocol
+    lr_scheduler: LRSchedulerProtocol
+    gradient_clipper: GradientClipper
+    exporter: ModelStageExporter
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            **super().state_dict(),
+            "logger": self.logger.state_dict(),
+            "task": self.task.state_dict(),
+            "metrics": self.metrics.state_dict(),
+            "optimizer": self.optimizer.state_dict(),
+            "lr_scheduler": self.lr_scheduler.state_dict(),
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        super().load_state_dict(state_dict)
+
+        self.logger.load_state_dict(state_dict["logger"])
+        self.task.load_state_dict(state_dict["task"])
+        self.metrics.load_state_dict(state_dict["metrics"])
+        self.optimizer.load_state_dict(state_dict["optimizer"])
+        self.lr_scheduler.load_state_dict(state_dict["lr_scheduler"])
+
+
+@dataclasses.dataclass(kw_only=True)
+class InferJobState(JobState):
+    """
+    Container for the state of an inference job.
+
+    Attributes:
+        task: The specific inference task logic definition.
+    """
+    task: InferenceTask
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            **super().state_dict(),
+            "task": self.task.state_dict(),
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        super().load_state_dict(state_dict)
+        self.task.load_state_dict(state_dict["task"])

d9d/lr_scheduler/__init__.py

@@ -0,0 +1,9 @@
+"""
+Utilities for learning rate scheduling.
+"""
+
+from .visualizer import visualize_lr_scheduler
+
+__all__ = [
+    "visualize_lr_scheduler"
+]

d9d/lr_scheduler/piecewise/__init__.py

@@ -0,0 +1,18 @@
+"""
+Implements flexible piecewise learning rate schedules via a builder pattern.
+"""
+
+from .builder import piecewise_schedule
+from .config import PiecewiseSchedulerConfig, piecewise_scheduler_from_config
+from .curves import CurveBase, CurveCosine, CurveExponential, CurveLinear, CurvePoly
+
+__all__ = [
+    "CurveBase",
+    "CurveCosine",
+    "CurveExponential",
+    "CurveLinear",
+    "CurvePoly",
+    "PiecewiseSchedulerConfig",
+    "piecewise_schedule",
+    "piecewise_scheduler_from_config"
+]

d9d/lr_scheduler/piecewise/builder.py

@@ -0,0 +1,152 @@
+from typing import Self
+
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LambdaLR
+
+from d9d.core.protocol import LRSchedulerProtocol
+
+from .curves import CurveBase
+from .engine import PiecewiseScheduleEngine, SchedulePhase
+
+
+class PiecewiseScheduleBuilder:
+    """
+    Builder for constructing multiphase learning rate schedules.
+    """
+
+    def __init__(
+            self,
+            initial_multiplier: float,
+            total_steps: int | None
+    ):
+        """
+        Constructs a new PiecewiseScheduleBuilder.
+
+        Args:
+            initial_multiplier: The starting learning rate multiplier (usually 0.0 or 1.0).
+            total_steps: The total number of training steps. Required if using percentage-based methods.
+        """
+
+        self._phases: list[SchedulePhase] = []
+        self._total_steps = total_steps
+        self._last_end_step = 0
+        self._last_multiplier = initial_multiplier
+
+    def for_steps(self, steps: int, target_multiplier: float, curve: CurveBase) -> Self:
+        """
+        Adds a schedule phase lasting for a specific number of steps.
+
+        Args:
+            steps: Duration of this phase in steps.
+            target_multiplier: The value of the multiplier at the end of this phase.
+            curve: The interpolation curve to use for bridging the start and end values.
+
+        Returns:
+            The builder instance for chaining.
+        """
+
+        self._phases.append(SchedulePhase(
+            start_step=self._last_end_step,
+            end_step=self._last_end_step + steps,
+            curve=curve,
+            start_value=self._last_multiplier,
+            end_value=target_multiplier
+        ))
+
+        self._last_end_step += steps
+        self._last_multiplier = target_multiplier
+
+        return self
+
+    def until_percentage(self, p: float, target_multiplier: float, curve: CurveBase) -> Self:
+        """
+        Adds a schedule phase lasting until a specific percentage of total training steps is reached.
+
+        Args:
+            p: The target percentage (0.0 to 1.0) of total_steps where this phase ends.
+            target_multiplier: The value of the multiplier at the end of this phase.
+            curve: The interpolation curve to use.
+
+        Returns:
+            The builder instance for chaining.
+
+        Raises:
+            ValueError: If total_steps was not provided in constructor or if the target
+                percentage implies a step count earlier than the current cursor.
+        """
+
+        if self._total_steps is None:
+            raise ValueError(
+                "You must define 'total_steps' in the constructor to use percentage-based methods."
+            )
+
+        if not 0.0 <= p <= 1.0:
+            raise ValueError("Percentage should be in range of [0.0, 1.0]")
+
+        target_step_abs = int(self._total_steps * p)
+        duration = target_step_abs - self._last_end_step
+
+        if duration < 0:
+            raise ValueError(
+                f"Target percentage {p} (step {target_step_abs}) is behind "
+                f"current cursor (step {self._last_end_step})."
+            )
+
+        return self.for_steps(duration, target_multiplier, curve)
+
+    def fill_rest(self, target_multiplier: float, curve: CurveBase) -> Self:
+        """
+        Adds a schedule phase that lasts from the current cursor until the end of training.
+
+        Args:
+            target_multiplier: The value of the multiplier at the very end of training.
+            curve: The interpolation curve to use.
+
+        Returns:
+            The builder instance for chaining.
+        """
+
+        return self.until_percentage(1.0, target_multiplier, curve)
+
+    def build(self, optimizer: Optimizer) -> LRSchedulerProtocol:
+        """
+        Finalizes the schedule and returns a PyTorch LR Scheduler.
+
+        Args:
+            optimizer: The optimizer to wrap.
+
+        Returns:
+            A scheduler configured with the defined phases.
+
+        Raises:
+            ValueError: If the defined phases exceed the total_steps provided.
+        """
+
+        if self._total_steps is not None and self._last_end_step > self._total_steps:
+            raise ValueError(
+                f"Schedule defined for {self._last_end_step} steps, but total_steps is {self._total_steps}."
+            )
+
+        engine = PiecewiseScheduleEngine(self._phases)
+        return LambdaLR(optimizer, engine.get_factor)
+
+
+def piecewise_schedule(
+        initial_multiplier: float,
+        total_steps: int | None = None
+) -> PiecewiseScheduleBuilder:
+    """
+    Entry point for creating a piecewise learning rate schedule.
+
+    Args:
+        initial_multiplier: The initial learning rate multiplier.
+        total_steps: Total training steps. Required for percentage-based scheduling.
+
+    Returns:
+        A builder instance to configure phases.
+    """
+
+    return PiecewiseScheduleBuilder(
+        initial_multiplier=initial_multiplier,
+        total_steps=total_steps
+    )