nemo_automodel.components.loss.utils

Module Contents

Functions

Name	Description
`_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.
`calculate_loss`	Calculate a logit-based or fused linear cross-entropy loss.

API

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