nemo_automodel.shared.tied_weights

Module Contents

Classes

Name	Description
`TieSupport`	Which `tie_word_embeddings` settings a model class supports.

Functions

Name	Description
`_get_module_by_normalized_name`	Return a module by FQN after applying wrapper-prefix normalization.
`_normalize_param_name`	Strip wrapper-specific prefixes from a parameter name.
`_same_tensor_storage`	Return whether two tensors alias the same local storage.
`ensure_tied_lm_head`	Ensure a local tied LM head actually aliases the input embedding.
`get_controlling_tie_word_embeddings`	Resolve the `tie_word_embeddings` flag that actually controls LM-head tying.
`get_input_embeddings_weight_and_name`	Return the input embedding weight and normalized name if present.
`get_lm_head_weight_and_name`	Return the first `lm_head.weight` parameter found on a model.
`has_local_tied_lm_head`	Return whether the current model partition has an actual tied LM head.
`is_tied_word_embeddings`	Check whether the model’s word embeddings are tied.

API

class nemo_automodel.shared.tied_weights.TieSupport

Bases: enum.Enum

Which tie_word_embeddings settings a model class supports.

Declared as the tie_word_embeddings_support class attribute on every registered model class and consulted by the model construction guard.

BOTH

= 'both'

TIED_ONLY

= 'tied_only'

UNTIED_ONLY

= 'untied_only'

nemo_automodel.shared.tied_weights._get_module_by_normalized_name(
    model: torch.nn.Module,
    normalized_module_name: str
) -> torch.nn.Module | None

Return a module by FQN after applying wrapper-prefix normalization.

nemo_automodel.shared.tied_weights._normalize_param_name(
    name: str
) -> str

Strip wrapper-specific prefixes from a parameter name.

nemo_automodel.shared.tied_weights._same_tensor_storage(
    left: torch.Tensor,
    right: torch.Tensor
) -> bool

Return whether two tensors alias the same local storage.

Parameters:

left

torch.Tensor

Tensor of arbitrary shape, or a DTensor whose local tensor should be compared.

right

torch.Tensor

Tensor with the same logical role as left. Its shape is not constrained here because callers validate shape compatibility.

Returns: bool

True when both local tensors have the same storage pointer and offset.

nemo_automodel.shared.tied_weights.ensure_tied_lm_head(
    model: torch.nn.Module
) -> bool

Ensure a local tied LM head actually aliases the input embedding.

Hugging Face tie_weights() is attempted first so model-specific tying rules remain authoritative. Direct assignment is the fallback for wrapped models where the generic method no longer reaches the local pair.

Parameters:

model

nn.Module

Model or pipeline stage to inspect and update.

Returns: bool

True if the local lm_head and input embedding are tied after the

nemo_automodel.shared.tied_weights.get_controlling_tie_word_embeddings(
    config: object
) -> bool

Resolve the tie_word_embeddings flag that actually controls LM-head tying.

Hugging Face ties lm_head based on the top-level config flag. Full Omni wrapper configs do not expose that flag directly, so they fall back to thinker_config. Other wrappers without a top-level flag fall back to text_config.

Parameters:

config

object

Model config exposing the relevant tying flag.

Returns: bool

The controlling tie_word_embeddings value.

nemo_automodel.shared.tied_weights.get_input_embeddings_weight_and_name(
    model: torch.nn.Module
) -> tuple[torch.Tensor | None, str | None]

Return the input embedding weight and normalized name if present.

Parameters:

model

nn.Module

Model to inspect.

Returns: torch.Tensor | None

The embedding weight tensor and its normalized FQN, or (None, None)

nemo_automodel.shared.tied_weights.get_lm_head_weight_and_name(
    model: torch.nn.Module
) -> tuple[torch.Tensor | None, str | None]

Return the first lm_head.weight parameter found on a model.

Parameters:

model

nn.Module

Model to inspect.

Returns: torch.Tensor | None

The parameter tensor and its normalized FQN, or (None, None) when

nemo_automodel.shared.tied_weights.has_local_tied_lm_head(
    model: torch.nn.Module
) -> bool

Return whether the current model partition has an actual tied LM head.

This is stricter than :func:is_tied_word_embeddings: pipeline stages often retain a tied config flag without owning both tensors, and speculative draft models can intentionally use different vocabulary sizes.

Parameters:

model

nn.Module

Model or pipeline stage to inspect.

Returns: bool

True when the local lm_head and input embedding both exist, have

nemo_automodel.shared.tied_weights.is_tied_word_embeddings(
    model: torch.nn.Module
) -> bool

Check whether the model’s word embeddings are tied.

A one-direction :class:TieSupport policy is authoritative. Only BOTH models and undeclared Hugging Face models need their config flag resolved.

Parameters:

model

nn.Module

Model to inspect.

Returns: bool

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