nemo_automodel.components.loss.utils#

Module Contents#

Functions#

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

_get_final_hidden_states

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

calculate_loss

Calculate the loss.

API#

nemo_automodel.components.loss.utils._get_lm_head_module(
model: torch.nn.Module,
) → 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._get_final_hidden_states(
model_output: Any,
) → Optional[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.calculate_loss(loss_fn, **kwargs) → torch.Tensor#

Calculate the loss.

Parameters:

  • loss_fn – Loss function.

  • **kwargs – Keyword arguments for the loss function.

Returns:

The loss.