nemo_automodel.components.optim.precision_warnings

Module Contents

Functions

Name	Description
`_has_trainable_bf16_param`	-
`_is_rank_zero`	-
`_is_torch_adam_config`	-
`_is_torch_adam_optimizer`	-
`_iter_optimizer_params`	-
`warn_if_torch_adam_with_bf16_params`	Warn about full-parameter bf16 training with vanilla torch Adam optimizers.

Data

_TORCH_ADAM_TARGETS

_WARNED_CONTEXTS

API

nemo_automodel.components.optim.precision_warnings._has_trainable_bf16_param(
    parameters: collections.abc.Iterable[torch.nn.Parameter]
) -> bool

nemo_automodel.components.optim.precision_warnings._is_rank_zero() -> bool

nemo_automodel.components.optim.precision_warnings._is_torch_adam_config(
    optimizer_cfg: typing.Any | None
) -> bool

nemo_automodel.components.optim.precision_warnings._is_torch_adam_optimizer(
    optimizer: torch.optim.Optimizer | collections.abc.Iterable[torch.optim.Optimizer] | None
) -> bool

nemo_automodel.components.optim.precision_warnings._iter_optimizer_params(
    optimizer: torch.optim.Optimizer | collections.abc.Iterable[torch.optim.Optimizer] | None
) -> collections.abc.Iterable[torch.nn.Parameter]

nemo_automodel.components.optim.precision_warnings.warn_if_torch_adam_with_bf16_params(
    optimizer: torch.optim.Optimizer | collections.abc.Iterable[torch.optim.Optimizer] | None = None,
    optimizer_cfg: typing.Any | None = None,
    parameters: collections.abc.Iterable[torch.nn.Parameter] | None = None,
    is_peft: bool = False,
    context: str = 'recipe',
    logger: logging.Logger | None = None
) -> None

Warn about full-parameter bf16 training with vanilla torch Adam optimizers.

nemo_automodel.components.optim.precision_warnings._TORCH_ADAM_TARGETS = {'torch.optim.Adam', 'torch.optim.AdamW', 'torch.optim.adam.Adam', 'torch.optim....

nemo_automodel.components.optim.precision_warnings._WARNED_CONTEXTS: set[str] = set()