d9d 0.1.0__py3-none-any.whl

d9d/lr_scheduler/piecewise/config.py

@@ -0,0 +1,176 @@
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field, PositiveInt
+from torch.optim import Optimizer
+
+from d9d.core.protocol import LRSchedulerProtocol
+
+from .builder import piecewise_schedule
+from .curves import CurveBase, CurveCosine, CurveExponential, CurveLinear, CurvePoly
+
+
+class CurveLinearConfig(BaseModel):
+    """
+    Configuration for linear interpolation.
+    """
+
+    type: Literal["linear"] = "linear"
+
+
+class CurveCosineConfig(BaseModel):
+    """
+    Configuration for cosine interpolation.
+    """
+
+    type: Literal["cosine"] = "cosine"
+
+
+class CurveExponentialConfig(BaseModel):
+    """
+    Configuration for exponential interpolation.
+    """
+
+    type: Literal["exponential"] = "exponential"
+
+
+class CurvePolyConfig(BaseModel):
+    """
+    Configuration for polynomial interpolation.
+
+    Attributes:
+        power: The exponent of the polynomial function.
+    """
+
+    type: Literal["poly"] = "poly"
+    power: float = 2.0
+
+
+AnyCurveConfig = Annotated[
+    CurveLinearConfig | CurveCosineConfig | CurveExponentialConfig | CurvePolyConfig,
+    Field(discriminator="type")
+]
+
+
+def curve_from_config(config: AnyCurveConfig) -> CurveBase:
+    """
+    Instantiates a concrete curve object from its configuration.
+
+    Args:
+        config: The configuration object.
+
+    Returns:
+        The instantiated curve.
+    """
+
+    match config:
+        case CurveLinearConfig():
+            return CurveLinear()
+        case CurvePolyConfig():
+            return CurvePoly(config.power)
+        case CurveExponentialConfig():
+            return CurveExponential()
+        case CurveCosineConfig():
+            return CurveCosine()
+
+
+class StepPhaseConfig(BaseModel):
+    """
+    Configuration for a phase defined by a fixed number of steps.
+
+    Attributes:
+        mode: Discriminator field, must be "steps".
+        steps: The absolute duration of this phase in steps.
+        target_multiplier: The multiplier value at the end of this phase.
+        curve: The interpolation curve configuration.
+    """
+
+    mode: Literal["steps"] = "steps"
+
+    steps: PositiveInt
+    target_multiplier: float
+    curve: AnyCurveConfig
+
+
+class PercentagePhaseConfig(BaseModel):
+    """
+    Configuration for a phase that lasts until a specific percentage of training is complete.
+
+    Attributes:
+        mode: Discriminator field, must be "percentage".
+        percentage: The target progress (0.0 to 1.0) where this phase ends.
+        target_multiplier: The multiplier value at the end of this phase.
+        curve: The interpolation curve configuration.
+    """
+
+    mode: Literal["percentage"] = "percentage"
+
+    percentage: float = Field(..., ge=0.0, le=1.0)
+    target_multiplier: float
+    curve: AnyCurveConfig
+
+
+class RestPhaseConfig(BaseModel):
+    """
+    Configuration for a phase that fills the remainder of the training duration.
+
+    Attributes:
+        mode: Discriminator field, must be "rest".
+        target_multiplier: The multiplier value at the very end of training.
+        curve: The interpolation curve configuration.
+    """
+
+    mode: Literal["rest"] = "rest"
+
+    target_multiplier: float
+    curve: AnyCurveConfig
+
+
+PhaseConfig = Annotated[
+    StepPhaseConfig | PercentagePhaseConfig | RestPhaseConfig,
+    Field(discriminator="mode")
+]
+
+
+class PiecewiseSchedulerConfig(BaseModel):
+    """
+    Declarative configuration for a piecewise learning rate scheduler.
+
+    Attributes:
+        initial_multiplier: The starting learning rate multiplier.
+        phases: A sequential list of phase configurations.
+    """
+
+    initial_multiplier: float
+    phases: list[PhaseConfig]
+
+
+def piecewise_scheduler_from_config(
+        config: PiecewiseSchedulerConfig,
+        optimizer: Optimizer,
+        total_steps: int | None
+) -> LRSchedulerProtocol:
+    """
+    Constructs a PyTorch scheduler from the provided configuration.
+
+    Args:
+        config: The scheduler configuration.
+        optimizer: The optimizer to wrap.
+        total_steps: The total number of training steps. Required if using percentage-based phases.
+
+    Returns:
+        A configured learning rate scheduler.
+    """
+
+    builder = piecewise_schedule(config.initial_multiplier, total_steps)
+
+    for phase in config.phases:
+        curve = curve_from_config(phase.curve)
+        match phase:
+            case StepPhaseConfig():
+                builder.for_steps(phase.steps, phase.target_multiplier, curve)
+            case PercentagePhaseConfig():
+                builder.until_percentage(phase.percentage, phase.target_multiplier, curve)
+            case RestPhaseConfig():
+                builder.fill_rest(phase.target_multiplier, curve)
+
+    return builder.build(optimizer)

d9d/lr_scheduler/piecewise/curves.py

@@ -0,0 +1,75 @@
+import abc
+import math
+
+
+class CurveBase(abc.ABC):
+    """
+    Abstract base class for interpolation curves used in scheduling.
+    """
+
+    @abc.abstractmethod
+    def compute(self, start: float, end: float, step_p: float) -> float:
+        """
+        Calculates the interpolated value.
+
+        Args:
+            start: The value at the beginning of the phase.
+            end: The value at the end of the phase.
+            step_p: Progress fraction through the phase (0.0 to 1.0).
+
+        Returns:
+            The interpolated value.
+        """
+
+
+class CurveLinear(CurveBase):
+    """
+    Linearly interpolates between start and end values.
+    """
+
+    def compute(self, start: float, end: float, step_p: float) -> float:
+        return start + (end - start) * step_p
+
+
+class CurveCosine(CurveBase):
+    """
+    Interpolates using a cosine annealing schedule (half-period cosine).
+    """
+
+    def compute(self, start: float, end: float, step_p: float) -> float:
+        cos_out = (1 + math.cos(math.pi * step_p)) / 2
+        return end + (start - end) * cos_out
+
+
+class CurvePoly(CurveBase):
+    """
+    Interpolates using a polynomial function.
+    """
+
+    def __init__(self, power: float):
+        """
+        Constructs a polynomial curve.
+
+        Args:
+            power: The exponent of the polynomial. 1.0 is linear, 2.0 is quadratic, etc.
+        """
+
+        self._power = power
+
+    def compute(self, start: float, end: float, step_p: float) -> float:
+        p_transformed = step_p ** self._power
+        return start + (end - start) * p_transformed
+
+
+class CurveExponential(CurveBase):
+    """
+    Interpolates exponentially between start and end values (log-space linear).
+    """
+
+    def compute(self, start: float, end: float, step_p: float) -> float:
+        eps = 1e-8
+        safe_start = max(start, eps)
+        safe_end = max(end, eps)
+
+        out_log = math.log(safe_start) + (math.log(safe_end) - math.log(safe_start)) * step_p
+        return math.exp(out_log)

d9d/lr_scheduler/piecewise/engine.py

@@ -0,0 +1,76 @@
+import dataclasses
+
+from .curves import CurveBase
+
+
+@dataclasses.dataclass
+class SchedulePhase:
+    """
+    Data container representing a single phase in a piecewise schedule.
+
+    Attributes:
+        start_step: The absolute step index where this phase begins.
+        end_step: The absolute step index where this phase ends.
+        start_value: The multiplier value at start_step.
+        end_value: The multiplier value at end_step.
+        curve: The interpolation logic for this phase.
+    """
+
+    start_step: int
+    end_step: int
+    start_value: float
+    end_value: float
+    curve: CurveBase
+
+
+class PiecewiseScheduleEngine:
+    """
+    Runtime engine that calculates multipliers based on a list of defined phases.
+    """
+
+    def __init__(self, phases: list[SchedulePhase]):
+        """
+        Constructs the schedule engine.
+
+        Args:
+            phases: A sequential list of schedule phases.
+
+        Raises:
+            ValueError: If the phases list is empty.
+        """
+
+        if len(phases) == 0:
+            raise ValueError("Scheduler should contain at least one phase")
+
+        self._phases = phases
+
+    def get_factor(self, step: int) -> float:
+        """
+        Computes the learning rate multiplier for the given step.
+
+        Args:
+            step: The global training step.
+
+        Returns:
+            The calculated multiplier. If the step is outside defined phases,
+            it clamps to the nearest boundary value.
+        """
+
+        if step < 0:
+            return self._phases[0].start_value
+
+        for phase in self._phases:
+            if not (phase.start_step <= step < phase.end_step):
+                continue
+
+            steps_in_phase = step - phase.start_step
+            phase_len = phase.end_step - phase.start_step
+            phase_progress = steps_in_phase / phase_len
+
+            return phase.curve.compute(
+                start=phase.start_value,
+                end=phase.end_value,
+                step_p=phase_progress
+            )
+
+        return self._phases[-1].end_value

d9d/lr_scheduler/visualizer.py

@@ -0,0 +1,74 @@
+from collections.abc import Callable
+
+from torch import nn
+from torch.optim import SGD, Optimizer
+from torch.optim.lr_scheduler import LRScheduler
+
+SchedulerFactory = Callable[[Optimizer], LRScheduler]
+
+
+def _get_history(factory: SchedulerFactory, num_steps: int, init_lr: float) -> list[float]:
+    optimizer = SGD(nn.Linear(1, 1).parameters(), lr=init_lr)
+
+    scheduler = factory(optimizer)
+
+    lrs = []
+
+    for _ in range(num_steps):
+        current_lr = optimizer.param_groups[0]["lr"]
+        lrs.append(current_lr)
+        scheduler.step()
+
+    return lrs
+
+
+def visualize_lr_scheduler(factory: SchedulerFactory, num_steps: int, init_lr: float = 1.0):
+    """
+    Visualizes the learning rate schedule using Plotly.
+
+    This function simulates the training process for `num_steps` to record the LR changes
+    and generates an interactive plot.
+
+    Args:
+        factory: A callable that accepts an Optimizer and returns an LRScheduler.
+        num_steps: The number of steps to simulate.
+        init_lr: The initial learning rate to set on the dummy optimizer.
+
+    Raises:
+        ImportError: If the `plotly` library is not installed.
+    """
+
+    try:
+        import plotly.graph_objects as go  # noqa: PLC0415
+    except ImportError as e:
+        raise ImportError("You have to install `plotly` dependency to use scheduler visualization") from e
+    lrs = _get_history(factory, num_steps, init_lr)
+    steps = list(range(num_steps))
+
+    fig = go.Figure()
+
+    fig.add_trace(go.Scatter(
+        x=steps,
+        y=lrs,
+        mode="lines",
+        name="Learning Rate",
+        line={"color": "#636EFA", "width": 3},
+        hovertemplate="<b>Step:</b> %{x}<br><b>LR:</b> %{y:.6f}<extra></extra>"
+    ))
+
+    fig.update_layout(
+        title={
+            "text": "Scheduler",
+            "y": 0.95,
+            "x": 0.5,
+            "xanchor": "center",
+            "yanchor": "top"
+        },
+        xaxis_title="Steps",
+        yaxis_title="Learning Rate",
+        template="plotly_white",
+        hovermode="x unified",
+        height=500
+    )
+
+    fig.show()

d9d/metric/__init__.py

@@ -0,0 +1,10 @@
+"""
+Distributed metric abstractions and implementations.
+"""
+
+
+from .abc import Metric
+
+__all__ = [
+    "Metric"
+]

d9d/metric/abc.py

@@ -0,0 +1,79 @@
+import abc
+from typing import Any, Generic, TypeVar
+
+import torch
+from torch.distributed.checkpoint.stateful import Stateful
+
+from d9d.core.dist_context import DistributedContext
+from d9d.core.types import TensorTree
+
+TComputeResult = TypeVar("TComputeResult", bound=TensorTree)
+
+
+class Metric(abc.ABC, Stateful, Generic[TComputeResult]):
+    """
+    Abstract base class for all metrics.
+
+    Metrics track statistics over time (e.g., during training) and can be synchronized
+    across distributed processes. They also support state persistence via the Stateful
+    interface.
+    """
+
+    @abc.abstractmethod
+    def update(self, *args: Any, **kwargs: Any):
+        """
+        Updates the metric state with a new batch of data.
+
+        Args:
+            *args: Positional arguments required by the specific metric implementation.
+            **kwargs: Keyword arguments required by the specific metric implementation.
+        """
+
+    @abc.abstractmethod
+    def trigger_sync(self, dist_context: DistributedContext):
+        """
+        Initiates the synchronization of the metric state across distributed processes.
+
+        This method should start the collective operations (e.g., all-reduce) required
+        to aggregate statistics, but should not block waiting for completion if possible.
+
+        Args:
+            dist_context: The distributed context.
+        """
+
+    @abc.abstractmethod
+    def wait_sync(self, dist_context: DistributedContext):
+        """
+        Waits for the synchronization initiated by `trigger_sync` to complete.
+
+        After this method returns, the metric state must be fully aggregated and
+        consistent across ranks.
+
+        Args:
+            dist_context: The distributed context.
+        """
+
+    @abc.abstractmethod
+    def compute(self) -> TComputeResult:
+        """
+        Computes the current value of the metric.
+
+        Returns:
+            The computed metric result (of type `TComputeResult`).
+                This can be a single `torch.Tensor` or `PyTree` structure (dict, list, etc.)
+                containing tensors, depending on how the subclass was typed.
+        """
+
+    @abc.abstractmethod
+    def reset(self):
+        """
+        Resets the internal state of the metric to the initial values.
+        """
+
+    def to(self, device: str | torch.device | int):
+        """
+        Moves a metric state to a specified device.
+
+        Args:
+            device: The device to move the metric state to.
+        """

d9d/metric/impl/__init__.py

@@ -0,0 +1,7 @@
+from .compose import ComposeMetric
+from .mean import WeightedMeanMetric
+
+__all__ = [
+    "ComposeMetric",
+    "WeightedMeanMetric"
+]

d9d/metric/impl/compose.py

@@ -0,0 +1,54 @@
+from collections.abc import Mapping
+from typing import Any
+
+import torch
+
+from d9d.core.dist_context import DistributedContext
+from d9d.metric import Metric
+
+
+class ComposeMetric(Metric[dict[str, Any]]):
+    def __init__(self, children: Mapping[str, Metric]):
+        self._children = children
+
+    def update(self, *args: Any, **kwargs: Any):
+        raise ValueError("Cannot update ComposeMetric directly - you can only update its children")
+
+    def __getitem__(self, item: str) -> Metric:
+        return self._children[item]
+
+    @property
+    def children(self) -> Mapping[str, Metric]:
+        return self._children
+
+    def trigger_sync(self, dist_context: DistributedContext):
+        for metric in self._children.values():
+            metric.trigger_sync(dist_context)
+
+    def wait_sync(self, dist_context: DistributedContext):
+        for metric in self._children.values():
+            metric.wait_sync(dist_context)
+
+    def compute(self) -> dict[str, Any]:
+        return {
+            metric_name: metric.compute()
+            for metric_name, metric in self._children.items()
+        }
+
+    def reset(self):
+        for metric in self._children.values():
+            metric.reset()
+
+    def to(self, device: str | torch.device | int):
+        for metric in self._children.values():
+            metric.to(device)
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            metric_name: metric.state_dict()
+            for metric_name, metric in self._children.items()
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        for metric_name, metric in self._children.items():
+            metric.load_state_dict(state_dict[metric_name])

d9d/metric/impl/mean.py

@@ -0,0 +1,94 @@
+from typing import Any
+
+import torch
+import torch.distributed as dist
+
+from d9d.core.dist_context import DistributedContext
+from d9d.metric import Metric
+
+
+class WeightedMeanMetric(Metric[torch.Tensor]):
+    """
+    Computes the weighted mean of values.
+
+    Tracks the sum of weighted values and the sum of weights.
+    """
+
+    def __init__(self):
+        """Constructs a WeightedMeanMetric object."""
+
+        super().__init__()
+        self._value = torch.scalar_tensor(0, dtype=torch.float32)
+        self._weight = torch.scalar_tensor(0, dtype=torch.float32)
+
+        self._is_synced = False
+        self._synced_value = torch.scalar_tensor(0, dtype=torch.float32)
+        self._synced_weight = torch.scalar_tensor(0, dtype=torch.float32)
+
+        self._handles: list[dist.Work] | None = None
+
+    def update(self, values: torch.Tensor, weights: torch.Tensor):
+        self._value += (values * weights).sum()
+        self._weight += weights.sum()
+
+        self._is_synced = False
+
+    def trigger_sync(self, dist_context: DistributedContext):
+        self._synced_value = self._value.clone()
+        self._synced_weight = self._weight.clone()
+        self._is_synced = True
+
+        self._handles = [
+            dist.all_reduce(self._synced_value, op=dist.ReduceOp.SUM, async_op=True),
+            dist.all_reduce(self._synced_weight, op=dist.ReduceOp.SUM, async_op=True)
+        ]
+
+    def wait_sync(self, dist_context: DistributedContext):
+        if self._handles is None:
+            raise RuntimeError("Sync was not triggered before")
+
+        for handle in self._handles:
+            handle.wait()
+        self._handles = None
+
+    def compute(self) -> torch.Tensor:
+        if self._is_synced:
+            return self._synced_value / self._synced_weight
+        else:
+            return self._value / self._weight
+
+    def reset(self):
+        self._value.fill_(0)
+        self._weight.fill_(0)
+        self._is_synced = False
+        self._handles = None
+
+    def to(self, device: str | torch.device | int):
+        self._weight = self._weight.to(device)
+        self._value = self._value.to(device)
+        self._synced_weight = self._synced_weight.to(device)
+        self._synced_value = self._synced_value.to(device)
+
+    @property
+    def accumulated_weight(self) -> torch.Tensor:
+        """
+        Returns the total weight accumulated so far.
+
+        Returns:
+            Scalar tensor with total weight.
+        """
+
+        if self._is_synced:
+            return self._synced_weight
+
+        return self._weight
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "value": self._value,
+            "weight": self._weight
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        self._value = state_dict["value"]
+        self._weight = state_dict["weight"]

d9d/model_state/io/__init__.py

@@ -0,0 +1,21 @@
+from .module_reader import load_model_state
+from .module_writer import (
+    save_model_state,
+    save_model_state_pipeline_parallel,
+)
+from .reader import read_model_state
+from .writer import (
+    write_model_state_distributed,
+    write_model_state_local,
+    write_model_state_pipeline_parallel,
+)
+
+__all__ = [
+    "load_model_state",
+    "read_model_state",
+    "save_model_state",
+    "save_model_state_pipeline_parallel",
+    "write_model_state_distributed",
+    "write_model_state_local",
+    "write_model_state_pipeline_parallel"
+]