llm.optimizers

llm/optimizers.py

Training optimizers.

get_optimizer

get_optimizer(
    name: Literal["lamb", "adam"],
    params: Iterable[Tensor] | Iterable[dict[str, Any]],
    lr: float,
    **kwargs: Any
) -> Optimizer

Get an optimizer by name.

Note

If ColossalAI is installed, fused versions of the optimizers will be created instead.

Parameters:

• name (Literal['lamb', 'adam']) –

  Name of the optimizer to load.

• params (Iterable[Tensor] | Iterable[dict[str, Any]]) –

  Parameters to be optimized.

• lr (float) –

  Learning rate.

• kwargs (Any, default: {} ) –

  Keyword arguments to pass to the optimizer.

Returns:

• Optimizer –

  Initialized optimizer.

Raises:

• ImportError –

  If name=='lamb' and ColossalAI is not installed.

• ValueError –

  If name is unknown.

Source code in llm/optimizers.py

def get_optimizer(
    name: Literal['lamb', 'adam'],
    params: Iterable[torch.Tensor] | Iterable[dict[str, Any]],
    lr: float,
    **kwargs: Any,
) -> torch.optim.Optimizer:
    """Get an optimizer by name.

    Note:
        If [ColossalAI](https://github.com/hpcaitech/ColossalAI){target=_blank}
        is installed, fused versions of the optimizers will be created instead.

    Args:
        name: Name of the optimizer to load.
        params: Parameters to be optimized.
        lr: Learning rate.
        kwargs: Keyword arguments to pass to the optimizer.

    Returns:
        Initialized optimizer.

    Raises:
        ImportError: If `name=='lamb'` and ColossalAI is not installed.
        ValueError: If `name` is unknown.
    """
    if name == 'adam':  # pragma: no cover
        if FUSED_IMPORT_ERROR is None:
            optimizer = FusedAdam(params, lr=lr, **kwargs)
        else:
            logger.warning(
                'ColossalAI with CUDA extensions is not installed so '
                'defaulting to native PyTorch Adam. Better performance can be '
                "enabled with ColossalAI's FusedAdam.",
            )
            optimizer = torch.optim.Adam(params, lr=lr, **kwargs)
    elif name == 'lamb':  # pragma: no cover
        if FUSED_IMPORT_ERROR is None:
            optimizer = FusedLAMB(params, lr=lr, **kwargs)
        else:
            raise ImportError(
                'FusedLamb is not available. ColossalAI with CUDA extensions '
                'is not installed.',
            ) from FUSED_IMPORT_ERROR
    else:
        raise ValueError(f'Unknown optimizer: {name}')

    return optimizer