`nemo_automodel.components.training.precision_warnings`#

Module Contents#

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

`_WARNED_CONTEXTS`
`_TORCH_ADAM_TARGETS`

nemo_automodel.components.training.precision_warnings._WARNED_CONTEXTS: set[str]#: ‘set(…)’

nemo_automodel.components.training.precision_warnings._TORCH_ADAM_TARGETS#: None

nemo_automodel.components.training.precision_warnings.warn_if_torch_adam_with_bf16_params( *, optimizer: torch.optim.Optimizer | collections.abc.Iterable[torch.optim.Optimizer] | None = None, optimizer_cfg: 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.training.precision_warnings._is_rank_zero() → bool#

nemo_automodel.components.training.precision_warnings._is_torch_adam_optimizer( optimizer: torch.optim.Optimizer | collections.abc.Iterable[torch.optim.Optimizer] | None, ) → bool#

nemo_automodel.components.training.precision_warnings._is_torch_adam_config(optimizer_cfg: Any | None) → bool#

nemo_automodel.components.training.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.training.precision_warnings._has_trainable_bf16_param( parameters: collections.abc.Iterable[torch.nn.Parameter], ) → bool#