> 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.components.training.embedding_row_repair

Repair near-zero input-embedding rows before optimizer construction.

## Module Contents

### Classes

| Name                                                                                                            | Description                                                         |
| --------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------- |
| [`EmbeddingRowRepairConfig`](#nemo_automodel-components-training-embedding_row_repair-EmbeddingRowRepairConfig) | Configuration for detecting and repairing damaged input embeddings. |
| [`EmbeddingRowRepairReport`](#nemo_automodel-components-training-embedding_row_repair-EmbeddingRowRepairReport) | Summary of a completed input-embedding repair.                      |
| [`_WeightView`](#nemo_automodel-components-training-embedding_row_repair-_WeightView)                           | Local tensor view and DTensor metadata for an embedding matrix.     |

### Functions

| Name                                                                                                                  | Description                                                              |
| --------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| [`_check_matching_layout`](#nemo_automodel-components-training-embedding_row_repair-_check_matching_layout)           | Validate that input and output embedding shards can be repaired rowwise. |
| [`_full_row_squared_norms`](#nemo_automodel-components-training-embedding_row_repair-_full_row_squared_norms)         | Compute complete row squared norms for the local vocabulary shard.       |
| [`_gather_unique_ids`](#nemo_automodel-components-training-embedding_row_repair-_gather_unique_ids)                   | Gather sorted global vocabulary row IDs across all ranks.                |
| [`_global_ids`](#nemo_automodel-components-training-embedding_row_repair-_global_ids)                                 | Map local vocabulary-row selections to global token IDs.                 |
| [`_global_min`](#nemo_automodel-components-training-embedding_row_repair-_global_min)                                 | Compute the global minimum of a scalar tensor.                           |
| [`_healthy_rms_norm`](#nemo_automodel-components-training-embedding_row_repair-_healthy_rms_norm)                     | Compute the RMS norm of healthy input-embedding rows.                    |
| [`_weight_view`](#nemo_automodel-components-training-embedding_row_repair-_weight_view)                               | Create a local view for a regular tensor or DTensor embedding matrix.    |
| [`repair_input_embedding_rows`](#nemo_automodel-components-training-embedding_row_repair-repair_input_embedding_rows) | Detect and repair near-zero input-embedding rows.                        |

### Data

[`logger`](#nemo_automodel-components-training-embedding_row_repair-logger)

### API

```python
class nemo_automodel.components.training.embedding_row_repair.EmbeddingRowRepairConfig(
    enabled: bool = True,
    min_norm: float = 1e-06,
    max_rows: int = 256
)
```

Dataclass

Configuration for detecting and repairing damaged input embeddings.

A row is repaired when its L2 norm is non-finite or no larger than
`min_norm`. Its replacement uses the corresponding output-embedding
direction, scaled to the RMS norm of the checkpoint's healthy input rows.
This preserves a token-specific direction while restoring the magnitude
expected by the pretrained model.

The operation runs once after base-checkpoint loading and model sharding,
before optimizer construction. It supports regular tensors and DTensors
whose two matrix dimensions are sharded or replicated.

```python
nemo_automodel.components.training.embedding_row_repair.EmbeddingRowRepairConfig.__post_init__() -> None
```

```python
nemo_automodel.components.training.embedding_row_repair.EmbeddingRowRepairConfig.apply(
    model: torch.nn.Module
) -> nemo_automodel.components.training.embedding_row_repair.EmbeddingRowRepairReport | None
```

Repair damaged input-embedding rows in `model` when enabled.

```python
class nemo_automodel.components.training.embedding_row_repair.EmbeddingRowRepairReport(
    input_embedding_name: str,
    output_embedding_name: str | None,
    repaired_row_ids: tuple[int, ...],
    min_norm_before: float,
    target_norm: float | None
)
```

Dataclass

Summary of a completed input-embedding repair.

```python
class nemo_automodel.components.training.embedding_row_repair._WeightView(
    local: torch.Tensor,
    global_shape: tuple[int, ...],
    local_shape: tuple[int, ...],
    global_offset: tuple[int, ...],
    mesh: object | None,
    placements: tuple[object, ...]
)
```

Dataclass

Local tensor view and DTensor metadata for an embedding matrix.

```python
nemo_automodel.components.training.embedding_row_repair._check_matching_layout(
    input_view: nemo_automodel.components.training.embedding_row_repair._WeightView,
    output_view: nemo_automodel.components.training.embedding_row_repair._WeightView
) -> None
```

Validate that input and output embedding shards can be repaired rowwise.

**Parameters:**

Input embedding matrix view with global shape \[vocab, hidden] and local tensor shape
\[local\_vocab, local\_hidden].

Output embedding matrix view with global shape \[vocab, hidden] and local tensor shape
\[local\_vocab, local\_hidden].

```python
nemo_automodel.components.training.embedding_row_repair._full_row_squared_norms(
    view: nemo_automodel.components.training.embedding_row_repair._WeightView
) -> torch.Tensor
```

Compute complete row squared norms for the local vocabulary shard.

**Parameters:**

Embedding matrix view with global shape \[vocab, hidden] and local tensor shape
\[local\_vocab, local\_hidden].

**Returns:** `torch.Tensor`

Tensor of shape \[local\_vocab] containing each local vocabulary row's squared L2 norm over the full

```python
nemo_automodel.components.training.embedding_row_repair._gather_unique_ids(
    local_ids: list[int]
) -> tuple[int, ...]
```

Gather sorted global vocabulary row IDs across all ranks.

```python
nemo_automodel.components.training.embedding_row_repair._global_ids(
    view: nemo_automodel.components.training.embedding_row_repair._WeightView,
    local_mask: torch.Tensor
) -> list[int]
```

Map local vocabulary-row selections to global token IDs.

**Parameters:**

Embedding matrix view with global shape \[vocab, hidden] and local tensor shape
\[local\_vocab, local\_hidden].

Boolean tensor of shape \[local\_vocab] selecting rows from `view.local`.

**Returns:** `list[int]`

Global vocabulary row IDs selected by `local_mask`.

```python
nemo_automodel.components.training.embedding_row_repair._global_min(
    value: torch.Tensor
) -> float
```

Compute the global minimum of a scalar tensor.

**Parameters:**

Scalar tensor of shape \[] containing this rank's local minimum.

**Returns:** `float`

Minimum scalar value across all distributed ranks.

```python
nemo_automodel.components.training.embedding_row_repair._healthy_rms_norm(
    row_squared_norms: torch.Tensor,
    healthy_mask: torch.Tensor
) -> float
```

Compute the RMS norm of healthy input-embedding rows.

**Parameters:**

Tensor of shape \[local\_vocab] containing each local vocabulary row's squared L2
norm over the full hidden axis.

Boolean tensor of shape \[local\_vocab] selecting healthy rows from `row_squared_norms`.

**Returns:** `float`

RMS L2 norm over all healthy global vocabulary rows.

```python
nemo_automodel.components.training.embedding_row_repair._weight_view(
    weight: torch.Tensor,
    name: str
) -> nemo_automodel.components.training.embedding_row_repair._WeightView
```

Create a local view for a regular tensor or DTensor embedding matrix.

**Parameters:**

Embedding matrix of shape \[vocab, hidden]. For DTensor inputs, the global shape is
\[vocab, hidden] and each placement must replicate or shard axis 0 (vocab) or axis 1 (hidden).

Parameter name used in validation errors.

**Returns:** `_WeightView`

View whose `local` tensor has shape \[local\_vocab, local\_hidden] and whose shape and offset metadata

```python
nemo_automodel.components.training.embedding_row_repair.repair_input_embedding_rows(
    model: torch.nn.Module,
    min_norm: float = 1e-06,
    max_rows: int = 256
) -> nemo_automodel.components.training.embedding_row_repair.EmbeddingRowRepairReport
```

Detect and repair near-zero input-embedding rows.

**Parameters:**

Loaded and optionally DTensor-sharded causal language model whose input and output embedding
weights have shape \[vocab, hidden].

Inclusive L2-norm threshold identifying a damaged row.

Safety bound; abort instead of mutating a broadly damaged checkpoint.

**Returns:** `EmbeddingRowRepairReport`

A report containing every repaired global vocabulary row ID.

```python
nemo_automodel.components.training.embedding_row_repair.logger = logging.getLogger(__name__)
```