loss

Classes

ChunkedCrossEntropy	Memory-bounded CE for replicated or vocabulary-TP logits.
KDLoss	TP-aware teacher-to-student Kullback-Leibler divergence.
TVDLoss	TP-aware total-variation distance between teacher and student.

Functions

build_distribution_loss

class ChunkedCrossEntropy

Bases: Module

Memory-bounded CE for replicated or vocabulary-TP logits.

Unlike AutoModel’s generic masked CE, a vocabulary-sharded DTensor remains sharded throughout the calculation. Token chunks bound the fp32 softmax workspace, and checkpointing recomputes that workspace during backward instead of retaining one vocabulary-sized buffer per token.

__init__(fp32_upcast=True, ignore_index=-100, reduction='sum', chunk_size=0, checkpoint_chunks=False)

Parameters:

• fp32_upcast (bool)

• ignore_index (int)

• reduction (str)

• chunk_size (int)

• checkpoint_chunks (bool)

forward(logits, labels, mask=None, num_label_tokens=None)

forward_no_checkpoint(logits, labels, num_label_tokens=None)

class KDLoss

Bases: _DistributionLoss

TP-aware teacher-to-student Kullback-Leibler divergence.

metric: Literal['kld', 'tvd'] = 'kld'

class TVDLoss

Bases: _DistributionLoss

TP-aware total-variation distance between teacher and student.

metric: Literal['kld', 'tvd'] = 'tvd'

build_distribution_loss(metric, **kwargs)

Parameters:

metric (str)

Return type:

_DistributionLoss