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, materializing DTensor weights when needed.
`calculate_loss`	Calculate the loss.

API

nemo_automodel.components.loss.utils._get_final_hidden_states(
    model_output: typing.Any
) -> typing.Optional[typing.Any]

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
) -> typing.Optional[torch.nn.Module]

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, materializing DTensor weights when needed.

nemo_automodel.components.loss.utils.calculate_loss(
    loss_fn,
    kwargs = {}
) -> torch.Tensor

Calculate the loss.

Parameters:

loss_fn

Loss function.

**kwargs

Defaults to {}

Keyword arguments for the loss function.

Returns: torch.Tensor

The loss.