Implementation:NVIDIA TransformerEngine Cross Entropy

Field	Value
Sources	TransformerEngine
Domains	Deep_Learning, PyTorch, Distributed
Last Updated	2026-02-07 14:00 GMT

Overview

Public API for cross entropy loss with Triton-fused computation, optional distributed vocabulary parallelism, label smoothing, and CUDA graph compatibility.

Description

This module provides the parallel_cross_entropy function and the underlying CrossEntropyFunction autograd function. The implementation delegates to Triton kernels for fused online-softmax + cross-entropy computation. Key features:

• Loss and gradient computation in FP32 regardless of input dtype (BF16/FP32 input supported)
• Distributed vocabulary parallelism via dist_process_group
• Label smoothing support
• CUDA graph capturable mode
• Backward compatibility with deprecated _input parameter

Usage

Use parallel_cross_entropy as a drop-in replacement for torch.nn.functional.cross_entropy when training with FP8 or distributed vocabulary parallelism.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/pytorch/cross_entropy.py

Lines

1--152

Signature

def parallel_cross_entropy(
    inp: torch.Tensor,
    target: torch.Tensor,
    label_smoothing: float = 0.0,
    reduce_loss: bool = False,
    dist_process_group: Optional[torch.distributed.ProcessGroup] = None,
    ignore_idx: int = -100,
    is_cg_capturable: bool = False,
) -> torch.Tensor: ...

Import

from transformer_engine.pytorch.cross_entropy import parallel_cross_entropy

I/O Contract

Inputs

Name	Type	Required	Description
inp	torch.Tensor	Yes	Input logits of shape (B, SQ, V) or (SQ, B, V)
target	torch.Tensor	Yes	Target indices of shape (B, SQ) or (SQ, B)
label_smoothing	float	No	Smoothing factor (default 0.0)
reduce_loss	bool	No	Return averaged scalar loss (default False)
dist_process_group	ProcessGroup	No	Distributed group for vocab-parallel loss
ignore_idx	int	No	Index to ignore (default -100)
is_cg_capturable	bool	No	CUDA graph capturable mode (default False)

Outputs

Name	Type	Description
loss	torch.Tensor	Cross entropy loss in FP32

Usage Examples

from transformer_engine.pytorch.cross_entropy import parallel_cross_entropy

loss = parallel_cross_entropy(
    inp=logits,        # (batch, seq_len, vocab_size)
    target=labels,     # (batch, seq_len)
    label_smoothing=0.1,
    reduce_loss=True,
    dist_process_group=tp_group,
)
loss.backward()

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements