nemo_automodel.components.models.common.tie_word_embeddings

Model-owned policy for input and output embedding ties.

Module Contents

Functions

Name	Description
`reject_tie_word_embeddings_flip`	Reject loading a checkpoint with `tie_word_embeddings` flipped from its own value.
`reject_unsupported_tie_word_embeddings`	Reject a `tie_word_embeddings` setting the model class cannot honor.

API

nemo_automodel.components.models.common.tie_word_embeddings.reject_tie_word_embeddings_flip(
    checkpoint_config: object,
    requested_config: object,
    model_class_name: str
) -> None

Reject loading a checkpoint with tie_word_embeddings flipped from its own value.

The class-level :class:TieSupport declaration cannot catch flipping the flag away from a specific checkpoint’s value (a BOTH class accepts either) — that is a (checkpoint, requested) property. NeMo AutoModel respects the checkpoint’s tie semantics, so a mismatch in either direction is rejected:

untied checkpoint requested tied -> would silently discard the trained lm_head;
tied checkpoint requested untied -> would leave a randomly-initialized lm_head (the adapters’ embed->head copy is gated on the flag).

Only applies to from_pretrained (there is a loaded checkpoint to compare against); from_config / scratch has no checkpoint, so the user’s config is authoritative.

Parameters:

checkpoint_config

object

The config parsed from the checkpoint, before user overrides.

requested_config

object

The config after user overrides are applied.

model_class_name

str

The resolved model class name to include in validation errors.

Raises:

NotImplementedError: If the controlling tie_word_embeddings flag differs between the checkpoint and the requested config.

nemo_automodel.components.models.common.tie_word_embeddings.reject_unsupported_tie_word_embeddings(
    model_cls: type,
    config: object
) -> None

Reject a tie_word_embeddings setting the model class cannot honor.

Reads the class’s declared :class:TieSupport policy from model_cls.tie_word_embeddings_support (defaulting to :attr:TieSupport.BOTH) and the controlling tie_word_embeddings flag from :func:get_controlling_tie_word_embeddings, then raises when the requested tying falls outside the supported set:

UNTIED_ONLY classes are separate-head architectures (HF default: distinct input/output embeddings) that never build a shared lm_head; honoring tie_word_embeddings=True would leave a randomly-initialized head.
TIED_ONLY classes share lm_head with the input embedding and ship checkpoints without a distinct lm_head.weight; honoring tie_word_embeddings=False would require materializing a head NeMo does not build and the checkpoint has no weights for.

Call at the top of __init__ on the original top-level config, before unwrapping to text_config / thinker_config or calling super().__init__(), so the correct controlling flag is inspected and construction fails fast.

Parameters:

model_cls

type

The constructing model class (type(self)). Its tie_word_embeddings_support attribute declares the policy, and its name is included in validation errors.

config

object

The model’s original (top-level) config.

Raises:

NotImplementedError: If tying is requested on a :attr:TieSupport.UNTIED_ONLY class, or untying is requested on a :attr:TieSupport.TIED_ONLY class.