nemo_automodel.components.loss.chunked_ce

Module Contents

Classes

Name	Description
`ChunkedCrossEntropy`	Chunk the LM-head projection and CE, recomputing each chunk in backward.
`_ChunkedLinearCE`	-

Functions

Name	Description
`_linear_cross_entropy`	Project one token chunk and compute CE.
`_validate_chunk_len`	Validate that `chunk_len` is positive.
`compute_cross_entropy`	Compute fp32 CE.

API

class nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy(
    chunk_len: int = 512,
    compile: bool = True,
    ignore_index: int = -100,
    reduction: str = 'sum'
)

Bases: LinearCrossEntropy

Chunk the LM-head projection and CE, recomputing each chunk in backward.

Unlike a logits-based loss, this consumes final hidden states and the projection weight. Neither full logits nor a full logits gradient is allocated. Only one chunk’s vocabulary activations are live at a time. Higher-order gradients are not supported.

_compute_loss

Callable[[torch.Tensor, torch.Tensor, torch.Tensor, int, torch.dtype | None], Tensor]

nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy.forward(
    hidden_states: torch.Tensor,
    labels: torch.Tensor,
    lm_weight: torch.Tensor,
    num_label_tokens: int | None = None,
    grad_reduce_group: torch.distributed.ProcessGroup | None = None,
    loss_weights: torch.Tensor | None = None,
    mask: torch.Tensor | None = None,
    logits_dtype: torch.dtype | None = None
) -> torch.Tensor

Compute CE directly from hidden states, with no full-logit allocation.

Parameters:

hidden_states

torch.Tensor

Rank-local states shaped […, hidden], with arbitrary leading token dimensions, including packed [tokens, hidden].

labels

torch.Tensor

Target indices shaped […] matching the token dimensions, or [tokens] for packed hidden states shaped [1, tokens, hidden].

lm_weight

torch.Tensor

Weight shaped [vocab, hidden]; FSDP DTensors use the LinearCrossEntropy.materialize_lm_weight contract. TP-sharded vocabulary/hidden states are not supported by this loss.

num_label_tokens

int | NoneDefaults to None

Global count normalizing a sum-reduced loss.

grad_reduce_group

dist.ProcessGroup | NoneDefaults to None

DP/CP group contributing independent token losses.

loss_weights

torch.Tensor | NoneDefaults to None

Optional constant multipliers shaped […], sum only.

mask

torch.Tensor | NoneDefaults to None

Optional mask shaped […]; zero positions are ignored.

logits_dtype

torch.dtype | NoneDefaults to None

Optional output cast after projection, matching the owning model. Projection uses lm_weight.dtype; mixed-dtype hidden states are converted per chunk, never the full weight.

Returns: torch.Tensor

FP32 scalar, or losses shaped […] for reduction=“none”. Inputs

nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy.prepare_lm_weight(
    lm_head: torch.nn.Module | None,
    model_config: transformers.PretrainedConfig | None = None,
    grad_reduce_group: torch.distributed.ProcessGroup | None = None
) -> torch.Tensor

Materialize one head in its effective projection dtype for all losses.

Parameters:

lm_head

nn.Module | None

Plain linear head with weight of global shape [vocab, hidden]. Its FSDP policy, when present, owns the compute dtype.

model_config

PretrainedConfig | NoneDefaults to None

Owning model config, used to reject post-projection transforms.

grad_reduce_group

dist.ProcessGroup | NoneDefaults to None

Group contributing independent token losses.

Returns: torch.Tensor

Dense [vocab, hidden] weight under materialize_lm_weight’s

nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy.validate_lm_head(
    lm_head: torch.nn.Module | None,
    model_config: transformers.PretrainedConfig | None = None
) -> torch.nn.Linear

staticmethod

Validate the head and the owning model’s output-logit configuration.

A linear module alone does not rule out post-projection transformations. Read the text configuration through the HF config contract, including multimodal models whose output settings live in a nested text config.

Parameters:

lm_head

nn.Module | None

Output head with weight of global shape [vocab, hidden].

model_config

PretrainedConfig | NoneDefaults to None

Owning model’s configuration, when available.

Returns: nn.Linear

The validated linear head, unchanged.

class nemo_automodel.components.loss.chunked_ce._ChunkedLinearCE()

Bases: Function

nemo_automodel.components.loss.chunked_ce._ChunkedLinearCE.backward(
    ctx: typing.Any,
    grad_out: torch.Tensor
) -> tuple[torch.Tensor | None, torch.Tensor | None, None, None, None, None, None]

staticmethod

Recompute one projection/CE graph at a time.

Parameters:

ctx

Any

Context holding [tokens, hidden] states, [vocab, hidden] weight and [tokens] labels.

grad_out

torch.Tensor

Upstream loss gradients shaped [tokens].

Returns: torch.Tensor | None

Hidden-state and weight gradients matching the original shapes and

nemo_automodel.components.loss.chunked_ce._ChunkedLinearCE.forward(
    ctx: typing.Any,
    hidden: torch.Tensor,
    weight: torch.Tensor,
    labels: torch.Tensor,
    chunk_len: int,
    ignore_index: int,
    logits_dtype: torch.dtype | None,
    compute_loss: collections.abc.Callable[[torch.Tensor, torch.Tensor, torch.Tensor, int, torch.dtype | None], torch.Tensor]
) -> torch.Tensor

staticmethod

Compute chunk losses without retaining logits.

Parameters:

ctx

Any

Autograd context.

hidden

torch.Tensor

Local hidden states shaped [tokens, hidden].

weight

torch.Tensor

Dense projection weight shaped [vocab, hidden].

labels

torch.Tensor

Target indices shaped [tokens].

chunk_len

int

Maximum token rows per projection.

ignore_index

int

Ignored target value.

logits_dtype

torch.dtype | None

Optional output cast matching the model projection.

compute_loss

Callable[[torch.Tensor, torch.Tensor, torch.Tensor, int, torch.dtype | None], torch.Tensor]

Eager or compiled per-chunk loss.

Returns: torch.Tensor

FP32 losses shaped [tokens]; inputs are not mutated.

nemo_automodel.components.loss.chunked_ce._linear_cross_entropy(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    labels: torch.Tensor,
    ignore_index: int,
    logits_dtype: torch.dtype | None
) -> torch.Tensor

Project one token chunk and compute CE.

Parameters:

hidden

torch.Tensor

Local hidden states shaped [chunk_tokens, hidden].

weight

torch.Tensor

Dense LM-head weight shaped [vocab, hidden].

labels

torch.Tensor

Target indices shaped [chunk_tokens].

ignore_index

int

Ignored target value.

logits_dtype

torch.dtype | None

Optional model-owned output dtype after projection.

Returns: torch.Tensor

FP32 losses shaped [chunk_tokens].

nemo_automodel.components.loss.chunked_ce._validate_chunk_len(
    chunk_len: int
) -> int

Validate that chunk_len is positive.

nemo_automodel.components.loss.chunked_ce.compute_cross_entropy(
    logits: torch.Tensor,
    targets: torch.Tensor,
    ignore_index: int = -100,
    reduction: str = 'sum'
) -> torch.Tensor

Compute fp32 CE.

Parameters:

logits

torch.Tensor

Scores shaped [tokens, vocab].

targets

torch.Tensor

Class indices shaped [tokens].

ignore_index

intDefaults to -100

Ignored target value.

reduction

strDefaults to 'sum'

PyTorch CE reduction.

Returns: torch.Tensor

Scalar loss, or [tokens] losses for reduction=“none”.