Function adamw_optimizer

axlearn/common/optimizers.py:741–793

AdamW optimizer with parameter scaling. N.B. The default weight-decay implementation is consistent with those in e.g. PyTorch & Optax, but inconsistent with the "decoupled" adamw weight decay formulation in Algorithm 2. To faithfully re

(
    learning_rate: schedule.Schedule,
    *,
    b1: float,
    b2: float,
    eps: float,
    weight_decay: float = 0,
    weight_decay_per_param_scale: Optional[Callable[[NestedOptParam], Any]] = None,
    mu_dtype: Optional[jnp.dtype] = None,
    adam_update_transformation: Optional[ConfigOr[PartitionedGradientTransformation]] = None,
)

def adamw_optimizer(
    learning_rate: schedule.Schedule,
    *,
    b1: float,
    b2: float,
    eps: float,
    weight_decay: float = 0,
    weight_decay_per_param_scale: Optional[Callable[[NestedOptParam], Any]] = None,
    mu_dtype: Optional[jnp.dtype] = None,
    adam_update_transformation: Optional[ConfigOr[PartitionedGradientTransformation]] = None,
) -> PartitionedGradientTransformation:
    """AdamW optimizer with parameter scaling.

    N.B. The default weight-decay implementation is consistent with
        those in e.g. PyTorch & Optax, but inconsistent with the "decoupled"
        adamw weight decay formulation in <https://arxiv.org/abs/1711.05101> Algorithm 2.
        To faithfully replicate Algorithm 2, use `adamw_decoupled_optimizer`.

    Args:
        learning_rate: the learning rate schedule.
        b1: the exponential decay rate for the 1st moment estimates.
        b2: the exponential decay rate for the 2nd moment estimates.
        eps: a small constant for numerical stability.
        weight_decay: optional rate at which to decay weights.
        weight_decay_per_param_scale: a Callable that returns a tree with same structure
            as the params PyTree, where each leaf is a float representing the per-param decay scale.
            The scale will be applied on top of the global decay rate:
            effective_decay_rate = global_decay_rate * per_param_scale.
            If None, all leaves will have a scale of 1.
        mu_dtype: optional `dtype` to be used for the first order accumulator;
            if `None` then the dtype is inferred from params and updates.
        adam_update_transformation: A transformation applied directly on the adam updates
            (but before weight decay). If None, no transformation is applied.

    Returns:
        A PartitionedGradientTransformation representing an AdamW optimizer with parameter scaling.
    """
    tx = [adam_partition(optax.scale_by_adam(b1=b1, b2=b2, eps=eps, mu_dtype=mu_dtype))]
    if adam_update_transformation is not None:
        tx.append(maybe_instantiate(adam_update_transformation))
    tx.extend(
        [
            add_decayed_weights(
                weight_decay=weight_decay,
                # Weight decay updates will already be scaled by the learning rate below.
                learning_rate_exponent=None,
                per_param_scale=weight_decay_per_param_scale,
            ),
            scale_by_schedule(scale_from_learning_rate(learning_rate)),
        ]
    )

    return chain(*tx)


def adamw_decoupled_optimizer(