colossalai.nn.loss.loss_1d

class colossalai.nn.loss.loss_1d.VocabParallelCrossEntropyLoss1D(reduction=True)

Vocab parallel cross entropy loss for 1D parallelism

Parameters

reduction (bool, optional) – whether to average the loss, defaults to True

forward(logits, targets)

Calculate loss between logits and targets

Parameters

• logits – Output logits of model

• targets – True targets from data