nemo_automodel.components.models.inkling2.loss

Fused softcapped cross entropy over the model’s usable vocabulary.

Module Contents

Classes

Name	Description
`Inkling2FusedLinearCrossEntropy`	Fused projection and cross entropy over the first `unpadded_vocab_size` head rows.

API

class nemo_automodel.components.models.inkling2.loss.Inkling2FusedLinearCrossEntropy(
    unpadded_vocab_size: int,
    logit_softcapping: float = 30.0,
    ignore_index: int = -100,
    reduction: str = 'sum'
)

Bases: FusedLinearCrossEntropy

Fused projection and cross entropy over the first unpadded_vocab_size head rows.

The kernel projects BF16 hidden states against the BF16 head rows with FP32 accumulation and applies the logit softcap and the cross entropy in FP32, which is the model’s output contract. Rows past the usable vocabulary never enter the softmax. An FP32 head is rounded to BF16, the cast FSDP applies to the model’s own projection.

nemo_automodel.components.models.inkling2.loss.Inkling2FusedLinearCrossEntropy.materialize_lm_weight(
    lm_weight: torch.Tensor,
    grad_reduce_group: torch.distributed.ProcessGroup | None = None
) -> torch.Tensor

Gather the head and keep its usable rows as the BF16 compute weight.

Parameters:

lm_weight

torch.Tensor

Head weight with global shape [padded_vocab, hidden], a regular tensor or an FSDP-sharded DTensor.

grad_reduce_group

dist.ProcessGroup | NoneDefaults to None

Ranks contributing independent token losses, as in the shared materialization contract.

Returns: torch.Tensor

BF16 tensor of shape [unpadded_vocab_size, hidden]. Gradients reach the