nemo_automodel.components.training.embedding_row_repair

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

Module Contents

Classes

Name	Description
`EmbeddingRowRepairConfig`	Configuration for detecting and repairing damaged input embeddings.
`EmbeddingRowRepairReport`	Summary of a completed input-embedding repair.
`_WeightView`	Local tensor view and DTensor metadata for an embedding matrix.

Functions

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

Data

logger

API

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.

enabled

bool = True

max_rows

int = 256

min_norm

float = 1e-06

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

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.

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.

input_embedding_name

str

min_norm_before

float

output_embedding_name

str | None

repaired_row_ids

tuple[int, ...]

target_norm

float | None

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.

global_offset

tuple[int, ...]

global_shape

tuple[int, ...]

local

Tensor

local_shape

tuple[int, ...]

mesh

object | None

placements

tuple[object, ...]

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_view

_WeightView

Input embedding matrix view with global shape [vocab, hidden] and local tensor shape [local_vocab, local_hidden].

output_view

_WeightView

Output embedding matrix view with global shape [vocab, hidden] and local tensor shape [local_vocab, local_hidden].

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:

view

_WeightView

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

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.

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:

view

_WeightView

Embedding matrix view with global shape [vocab, hidden] and local tensor shape [local_vocab, local_hidden].

local_mask

torch.Tensor

Boolean tensor of shape [local_vocab] selecting rows from view.local.

Returns: list[int]

Global vocabulary row IDs selected by local_mask.

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

Compute the global minimum of a scalar tensor.

Parameters:

value

torch.Tensor

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

Returns: float

Minimum scalar value across all distributed ranks.

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:

row_squared_norms

torch.Tensor

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

healthy_mask

torch.Tensor

Boolean tensor of shape [local_vocab] selecting healthy rows from row_squared_norms.

Returns: float

RMS L2 norm over all healthy global vocabulary rows.

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:

weight

torch.Tensor

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).

name

str

Parameter name used in validation errors.

Returns: _WeightView

View whose local tensor has shape [local_vocab, local_hidden] and whose shape and offset metadata

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:

model

nn.Module

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

min_norm

floatDefaults to 1e-06

Inclusive L2-norm threshold identifying a damaged row.

max_rows

intDefaults to 256

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

Returns: EmbeddingRowRepairReport

A report containing every repaired global vocabulary row ID.

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