Skip to content

Optimizer

About

The loop creates one optimizer per local model stage through an OptimizerProvider. You can configure a standard optimizer with AutoOptimizerProvider, or implement the OptimizerProvider protocol yourself.

Auto Optimizer

The d9d.loop.auto package builds standard optimizers from a Pydantic configuration. It supports AdamW, Adam, SGD and StochasticAdamW. The name field selects the optimizer.

Custom Optimizer

For a custom optimizer, implement the OptimizerProvider protocol. It receives the model stage and returns a PyTorch optimizer.

Usage

AutoOptimizerConfig is a discriminated union, so validate it with a Pydantic TypeAdapter:

1
2
3
4
5
6
from pydantic import TypeAdapter

from d9d.loop.auto import AutoOptimizerConfig, AutoOptimizerProvider

config = TypeAdapter(AutoOptimizerConfig).validate_json('{"name": "adamw", "lr": 1e-4}')
provider = AutoOptimizerProvider(config)

Inside a larger Pydantic config, declare a field of type AutoOptimizerConfig instead.

API Reference

d9d.loop.auto.auto_optimizer

Optimizer provider built from a configuration.

AutoOptimizerConfig = Annotated[StochasticAdamWOptimizerConfig | AdamWOptimizerConfig | AdamOptimizerConfig | SGDOptimizerConfig, Field(discriminator='name')] module-attribute

Discriminated union of the supported optimizer configurations.

AdamOptimizerConfig

Bases: BaseAutoOptimizerConfig

Configuration for the PyTorch Adam optimizer.

Attributes:

Name Type Description
name Literal['adam']

Discriminator field. Always "adam".
lr float

The learning rate.
betas tuple[float, float]

Coefficients for computing running averages of gradient and its square.
eps float

Term added to the denominator to improve numerical stability.
weight_decay float

Weight decay coefficient.
decoupled_weight_decay bool

Whether to apply decoupled weight decay.
amsgrad bool

Whether to use the AMSGrad variant.
maximize bool

Whether to maximize the objective instead of minimizing it.
build(params)

Builds fused Adam with the configured parameters.

Returns:

Type Description
Optimizer

The constructed Adam optimizer.

AdamWOptimizerConfig

Bases: BaseAutoOptimizerConfig

Configuration for the PyTorch AdamW optimizer.

Attributes:

Name Type Description
name Literal['adamw']

Discriminator field. Always "adamw".
lr float

The learning rate.
betas tuple[float, float]

Coefficients for computing running averages of gradient and its square.
eps float

Term added to the denominator to improve numerical stability.
weight_decay float

Weight decay coefficient.
amsgrad bool

Whether to use the AMSGrad variant.
maximize bool

Whether to maximize the objective instead of minimizing it.
build(params)

Builds fused AdamW with the configured parameters.

Returns:

Type Description
Optimizer

The constructed AdamW optimizer.

AutoOptimizerProvider

Bases: OptimizerProvider

Provider that builds a PyTorch optimizer from a configuration object.

__init__(config)

Constructs the AutoOptimizerProvider object.

BaseAutoOptimizerConfig

Bases: BaseModel, ABC

Abstract base class for optimizer configurations.

build(params) abstractmethod

Creates the PyTorch optimizer instance.

Parameters:

Name Type Description Default
params Iterable[Parameter]

An iterable of model parameters to optimize.

 required

Returns:

Type Description
Optimizer

The instantiated optimizer.

SGDOptimizerConfig

Bases: BaseAutoOptimizerConfig

Configuration for the PyTorch SGD optimizer.

Attributes:

Name Type Description
name Literal['sgd']

Discriminator field. Always "sgd".
lr float

The learning rate.
momentum float

Momentum factor.
dampening float

Dampening for momentum.
weight_decay float

Weight decay (L2 penalty).
nesterov bool

Enables Nesterov momentum.
maximize bool

Whether to maximize the objective instead of minimizing it.
build(params)

Builds fused SGD with the configured parameters.

Returns:

Type Description
Optimizer

The constructed SGD optimizer.

StochasticAdamWOptimizerConfig

Bases: BaseAutoOptimizerConfig

Configuration for the Stochastic AdamW optimizer.

Attributes:

Name Type Description
name Literal['stochastic_adamw']

Discriminator field. Always "stochastic_adamw".
lr float

Learning rate.
betas tuple[float, float]

Coefficients used for computing running averages of gradient and its square.
eps float

Term added to the denominator to improve numerical stability.
weight_decay float

Weight decay coefficient.
state_dtype str

The name of the torch dtype for the optimizer states, e.g. "bfloat16".
build(params)

Builds StochasticAdamW with the configured parameters.

Returns:

Type Description
Optimizer

The constructed StochasticAdamW optimizer.

d9d.loop.control.optimizer_provider

Interface for providing an optimizer.

InitializeOptimizerStageContext dataclass

Context data required to initialize an optimizer.

Attributes:

Name Type Description
dist_context DistributedContext

The distributed context.
model Module

The model instance for which parameters will be optimized.

OptimizerProvider

Bases: Protocol

Protocol for defining how optimizers are created for model pipeline stages.

__call__(context) abstractmethod

Initializes the optimizer for a specific training stage.

Parameters:

Name Type Description Default
context InitializeOptimizerStageContext

Context for this operation.

 required

Returns:

Type Description
Optimizer

The instantiated PyTorch optimizer.