nemo_automodel.shared.embedding_padding

DTensor-safe zeroing of one embedding row (the padding_idx step of weight init).

Module Contents

Functions

Name	Description
`zero_embedding_row_`	Zero `weight[row]` in place without integer-indexing a DTensor.

Data

_SUPPORTED_PLACEMENTS

logger

API

nemo_automodel.shared.embedding_padding.zero_embedding_row_(
    weight: torch.Tensor,
    row: int
) -> bool

Zero weight[row] in place without integer-indexing a DTensor.

weight[row].zero_() on a DTensor whose vocabulary dim is sharded triggers a redistribute (an all-gather of the whole embedding) and fails outright for TP shards. This touches only the rank-local shard, and only when that shard owns the row.

Parameters:

weight

torch.Tensor

Embedding matrix of shape [vocab, hidden]; a plain tensor or a DTensor whose placements are Replicate / Shard / _StridedShard (FSDP2 x TP) on either matrix axis.

row

int

Global row index to zero (negative indices count from the end).

Returns: bool

True when this rank held part of the row and zeroed it; False when the row lives

nemo_automodel.shared.embedding_padding._SUPPORTED_PLACEMENTS = (Replicate, Shard) + ((_StridedShard,) if _StridedShard else ())

nemo_automodel.shared.embedding_padding.logger = logging.getLogger(__name__)