nemo_automodel.components.loss.chunked_ce

Module Contents

Classes

Name	Description
`ChunkedCrossEntropy`	Cross-entropy loss computed over sequence chunks.
`_ChunkedCrossEntropySum`	Sum-reduced cross-entropy with recomputed fp32 chunk activations.

Functions

Name	Description
`_validate_chunk_len`	Validate that `chunk_len` is positive.
`compute_cross_entropy`	Computes the cross-entropy loss between logits and targets.

Data

_compiled_compute_cross_entropy

API

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

Bases: Module

Cross-entropy loss computed over sequence chunks.

chunk_len

= _validate_chunk_len(chunk_len)

nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy.forward(
    logits: torch.Tensor,
    labels: torch.Tensor,
    mask: typing.Optional[torch.Tensor] = None,
    num_label_tokens: typing.Optional[int] = None
) -> torch.Tensor

Computes cross-entropy loss in chunks to handle long sequences more efficiently.

Parameters:

logits

torch.Tensor

Model output tensor of shape […, vocab], with arbitrary leading token dimensions.

labels

torch.Tensor

Target indices of shape […] matching logits.shape[:-1]. When mask is provided, ignored positions are replaced with ignore_index in this tensor.

mask

Optional[torch.Tensor]Defaults to None

Optional tensor of shape […] matching labels. Nonzero positions contribute to the loss and zero positions are ignored.

num_label_tokens

Optional[int]Defaults to None

Optional global count used to normalize the sum-reduced scalar loss.

Returns: torch.Tensor

Scalar tensor containing the reduced cross-entropy loss.

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

Bases: Function

Sum-reduced cross-entropy with recomputed fp32 chunk activations.

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

staticmethod

Recompute the fp32 softmax chunks and return the logits gradient.

Parameters:

ctx

Autograd context containing the original-dtype tensors saved by :meth:forward.

grad_out

torch.Tensor

Scalar tensor containing the upstream loss gradient.

Returns: torch.Tensor

Tuple whose first element is the logits gradient of shape [tokens,

nemo_automodel.components.loss.chunked_ce._ChunkedCrossEntropySum.forward(
    ctx,
    logits: torch.Tensor,
    labels: torch.Tensor,
    ignore_index: int,
    chunk_len: int
) -> torch.Tensor

staticmethod

Compute sum-reduced cross-entropy one fp32 chunk at a time.

Parameters:

ctx

Autograd context used to save tensors for backward.

logits

torch.Tensor

Prediction scores of shape [tokens, vocab]. The tensor stays in its original floating-point dtype and is not mutated.

labels

torch.Tensor

Target class indices of shape [tokens]. Positions equal to ignore_index contribute zero loss and gradient.

ignore_index

int

Target value excluded from the loss.

chunk_len

int

Maximum number of token rows upcast to fp32 at once.

Returns: torch.Tensor

Scalar fp32 tensor containing the sum-reduced loss.

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 = -100,
    reduction = 'sum'
)

Computes the cross-entropy loss between logits and targets.

Parameters:

logits

torch.Tensor

Model predictions of shape (sequence_length, num_classes).

targets

torch.Tensor

Ground-truth labels of shape (sequence_length,).

ignore_index

intDefaults to -100

Target value that is ignored when computing the loss. Defaults to -100.

Returns:

torch.Tensor: The sum of cross-entropy losses over the sequence.

nemo_automodel.components.loss.chunked_ce._compiled_compute_cross_entropy = None