> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.shared.tied_weights

## Module Contents

### Classes

| Name                                                           | Description                                                  |
| -------------------------------------------------------------- | ------------------------------------------------------------ |
| [`TieSupport`](#nemo_automodel-shared-tied_weights-TieSupport) | Which `tie_word_embeddings` settings a model class supports. |

### Functions

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

### API

```python
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.

```python
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.

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

Strip wrapper-specific prefixes from a parameter name.

```python
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:**

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

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.

```python
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 or pipeline stage to inspect and update.

**Returns:** `bool`

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

```python
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:**

Model config exposing the relevant tying flag.

**Returns:** `bool`

The controlling `tie_word_embeddings` value.

```python
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 to inspect.

**Returns:** `torch.Tensor | None`

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

```python
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 to inspect.

**Returns:** `torch.Tensor | None`

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

```python
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 or pipeline stage to inspect.

**Returns:** `bool`

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

```python
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 to inspect.

**Returns:** `bool`

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