colossalai.nn.loss.loss_3d

class colossalai.nn.loss.loss_3d.CrossEntropyLoss3D(reduction=True, *args, **kwargs)

Cross entropy loss for 3D parallelism.

Parameters

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

The args and kwargs should include parameters below:

weight (Tensor, optional)
size_average (bool, optional)
ignore_index (int, optional)
reduce (bool, optional)
label_smoothing (float, optional)

More details about args, kwargs and torch.nn.functional.cross_entropy could be found in Cross_entropy.

forward(logits, targets)

Calculate loss between logits and targets.

Parameters

• logits (torch.tensor) – Predicted unnormalized scores (often referred to as logits).

• targets (torch.tensor) – Ground truth class indices or class probabilities.

class colossalai.nn.loss.loss_3d.VocabParallelCrossEntropyLoss3D(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 (torch.tensor) – Predicted unnormalized scores (often referred to as logits).

• targets (torch.tensor) – Ground truth class indices or class probabilities.