nemo_automodel.components.loss.linear_ce

Module Contents

Classes

Name	Description
`FusedLinearCrossEntropy`	Fused linear-projection and cross-entropy loss module.

Functions

Name	Description
`_get_triton_version`	-
`new_is_triton_greater_or_equal`	Check if pytorch-triton/triton version is greater than or equal to the specified version.
`new_is_triton_greater_or_equal_3_2_0`	Check if pytorch-triton/triton version is greater than or equal to 3.1.0.

Data

HAVE_CUT_CROSS_ENTROPY

API

class nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy(
    ignore_index: int = -100,
    logit_softcapping: float = 0,
    reduction: str = 'sum'
)

Bases: Module

Fused linear-projection and cross-entropy loss module.

nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy.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
) -> torch.Tensor

Compute fused linear cross entropy matching PyTorch behavior.

Parameters:

hidden_states

torch.Tensor

Rank-local hidden states with shape [batch, sequence, hidden].

labels

torch.Tensor

Rank-local target token IDs with shape [batch, sequence].

lm_weight

torch.Tensor

LM-head weight with global shape [vocab, hidden]. It may be a regular tensor or an FSDP-sharded DTensor.

num_label_tokens

int | NoneDefaults to None

Global number of non-padding target tokens used to normalize a sum-reduced loss.

grad_reduce_group

dist.ProcessGroup | NoneDefaults to None

Group that contributes independent loss shards when lm_weight is a sharded DTensor.

Returns: torch.Tensor

Scalar loss tensor on the same device as hidden_states. The

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

staticmethod

Materialize an LM-head DTensor with gradient-correct reduction semantics.

Fused linear CE consumes the LM-head weight outside the owning FSDP module’s forward. Each data/context-parallel rank therefore computes a rank-local full-weight gradient. A plain DTensor.full_tensor() marks that gradient as replicated, so backward only slices the local result into the owned shard instead of combining peer contributions.

Parameters:

lm_weight

torch.Tensor

LM-head weight with global shape [vocab, hidden]. A regular tensor is returned unchanged. A DTensor may have any FSDP sharding placement over its device mesh and is gathered to a rank-local regular tensor with the global shape, device, and dtype.

grad_reduce_group

dist.ProcessGroup | NoneDefaults to None

Process group whose ranks contribute independent token losses. Its size must match the LM-head DTensor mesh.

Returns: torch.Tensor

Regular tensor with shape [vocab, hidden]. For a DTensor input,

Raises:

ValueError: If a trainable sharded weight has no matching reduction group. This fails closed instead of producing rank-local shards.

nemo_automodel.components.loss.linear_ce._get_triton_version()

nemo_automodel.components.loss.linear_ce.new_is_triton_greater_or_equal(
    version_str
)

Check if pytorch-triton/triton version is greater than or equal to the specified version.

Parameters:

version_str

Version string to check

Returns:

True if pytorch-triton/triton version >= specified version

nemo_automodel.components.loss.linear_ce.new_is_triton_greater_or_equal_3_2_0()

Check if pytorch-triton/triton version is greater than or equal to 3.1.0.

Returns:

True if pytorch-triton/triton version >= 3.1.0

nemo_automodel.components.loss.linear_ce.HAVE_CUT_CROSS_ENTROPY = True