nemo_automodel.components.checkpoint.utils

Module Contents

Functions

resolve_trust_remote_code	Whitelist NVIDIA models to allow remote code execution.
is_tied_word_embeddings	Check if the model’s word embeddings are tied.
_get_checkpoint_tensor_dtypes	Inspect checkpoint tensors and return their exact dtypes by key.

API

nemo_automodel.components.checkpoint.utils.resolve_trust_remote_code(pretrained_model_name_or_path)

Whitelist NVIDIA models to allow remote code execution.

Parameters:

pretrained_model_name_or_path (str) – The name or path of the pretrained model.

Returns:

True if the model should be loaded with trust_remote_code, False otherwise.

Return type:

bool

nemo_automodel.components.checkpoint.utils.is_tied_word_embeddings(model: torch.nn.Module) → bool

Check if the model’s word embeddings are tied.

Parameters:

model (nn.Module) – The model to check.

Returns:

True if the model’s word embeddings are tied, False otherwise.

Return type:

bool

nemo_automodel.components.checkpoint.utils._get_checkpoint_tensor_dtypes(

pretrained_model_name_or_path: str,

hf_config: Any,

load_kwargs: collections.abc.Mapping[str, object] | None = None,

) → dict[str, torch.dtype]

Inspect checkpoint tensors and return their exact dtypes by key.

This reads checkpoint metadata only by loading tensors on the meta device, so it preserves the per-tensor dtype information without materializing full checkpoint weights in memory.