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.

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.

Returns

the loss between logits and targets.

Return type

float

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 (torch.tensor) – Predicted unnormalized scores (often referred to as logits).

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