colossalai.nn.loss.loss_2d

class colossalai.nn.loss.loss_2d.CrossEntropyLoss2D(reduction=True, *args, **kwargs)

Cross entropy loss for 2D parallelism

Parameters

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

• args – Args for loss function

• kwargs – Kwargs for loss function

forward(logits, targets)

Calculate loss between logits and targets

Parameters

• logits – Output logits of model

• targets – True targets from data

class colossalai.nn.loss.loss_2d.VocabParallelCrossEntropyLoss2D(reduction=True)

Vocab parallel cross entropy loss for 2D 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