Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo Automodelnemo_automodel.shared

nemo_automodel.shared

View as Markdown

Submodules

Package Contents

Functions

NameDescription
dtype_from_strTranslate a str val of a dtype into the corresponding torch.dtype
patch_t5_layer_normReplace apex’s FusedRMSNorm with a native T5LayerNorm in the T5 module.

API

nemo_automodel.shared.utils.dtype_from_str(
    val,
    default = torch.bfloat16
)

Translate a str val of a dtype into the corresponding torch.dtype Args: val (str): the dotted path of the dtype (e.g., “torch.bfloat16”).

Returns:

torch.dtype: the actual dtype (e.g., torch.bfloat16)

nemo_automodel.shared.transformers_patches.patch_t5_layer_norm() -> None

Replace apex’s FusedRMSNorm with a native T5LayerNorm in the T5 module.

Apex’s FusedRMSNorm doesn’t support bfloat16, but the native T5LayerNorm handles it correctly by upcasting to fp32 internally for numerical stability. This must be called before loading any T5 models.

This function is idempotent and safe to call multiple times.