nemo_automodel.components.loss.utils

Module Contents

Functions

Name	Description
`_count_label_tokens`	Count supervised entries in an integer target tensor of any shape.
`_get_final_hidden_states`	Return the final hidden-states tensor from an HF-like model output.
`_get_lm_head_module`	Return the model’s LM-head module, if one can be found.
`_get_lm_head_weight`	Return the model’s LM-head weight without changing its distributed layout.
`_get_loss_ignore_index`	Return the label sentinel consumed by `loss_fn`.
`_normalize_kd_labels`	Align KD supervision with the main loss mask.
`_normalize_loss_labels`	Map dataset padding to a loss’s configured ignore index.
`calculate_loss`	Calculate a logit-based or fused linear cross-entropy loss.

Data

_DATASET_IGNORE_INDEX

API

nemo_automodel.components.loss.utils._count_label_tokens(
    labels: torch.Tensor,
    ignore_index: int
) -> int

Count supervised entries in an integer target tensor of any shape.

Both the dataset’s -100 padding and the configured ignore_index are excluded. The return value is a Python integer.

nemo_automodel.components.loss.utils._get_final_hidden_states(
    model_output: typing.Any
) -> typing.Any | None

Return the final hidden-states tensor from an HF-like model output.

Local copy of components.training.model_output_utils.get_final_hidden_states to keep components/loss/ import-independent from components/training/.

nemo_automodel.components.loss.utils._get_lm_head_module(
    model: torch.nn.Module
) -> torch.nn.Module | None

Return the model’s LM-head module, if one can be found.

Local copy of components.utils.model_utils.get_lm_head_module to keep components/loss/ import-independent from components/utils/ (see the Components must not import each other import-linter contract).

nemo_automodel.components.loss.utils._get_lm_head_weight(
    model: torch.nn.Module
) -> torch.Tensor

Return the model’s LM-head weight without changing its distributed layout.

nemo_automodel.components.loss.utils._get_loss_ignore_index(
    loss_fn: object
) -> int

Return the label sentinel consumed by loss_fn.

nemo_automodel.components.loss.utils._normalize_kd_labels(
    labels: torch.Tensor,
    loss_ignore_index: int,
    kd_ignore_index: int
) -> torch.Tensor

Align KD supervision with the main loss mask.

Parameters:

labels

torch.Tensor

Integer target tensor of any shape.

loss_ignore_index

int

Label sentinel consumed by the main loss.

kd_ignore_index

int

Label sentinel consumed by the KD loss.

Returns: torch.Tensor

Target tensor with the same shape, dtype, and device as labels.

nemo_automodel.components.loss.utils._normalize_loss_labels(
    labels: torch.Tensor,
    ignore_index: int
) -> torch.Tensor

Map dataset padding to a loss’s configured ignore index.

Parameters:

labels

torch.Tensor

Integer target tensor of any shape.

ignore_index

int

Label sentinel consumed by the loss.

Returns: torch.Tensor

Target tensor with the same shape, dtype, and device as labels.

nemo_automodel.components.loss.utils.calculate_loss(
    loss_fn: torch.nn.Module,
    kwargs: typing.Any = {}
) -> torch.Tensor

Calculate a logit-based or fused linear cross-entropy loss.

Parameters:

loss_fn

nn.Module

Loss module. FusedLinearCrossEntropy consumes hidden_states with shape [batch, sequence, hidden], labels with shape [batch, sequence], and an LM-head weight with global shape [vocab, hidden]. Other loss modules consume logits with shape [batch, sequence, vocab] and labels.

**kwargs

AnyDefaults to {}

Loss inputs. Rank-local tensors keep their existing layout; grad_reduce_group describes the ranks contributing independent fused-loss shards. The caller’s mapping and tensors are not mutated.

Returns: torch.Tensor

Scalar loss tensor that does not alias an input.

nemo_automodel.components.loss.utils._DATASET_IGNORE_INDEX = -100