Is this page helpful?

core.tensor_parallel.cross_entropy#

Module Contents#

Classes#

VocabParallelCrossEntropy

Computes the Cross Entropy Loss splitting the Vocab size across tensor parallel ranks. This implementation is used in both fused and unfused cross entropy implementations

_VocabParallelCrossEntropy

Functions#

vocab_parallel_cross_entropy

Performs cross entropy loss when logits are split across tensor parallel ranks

unfused_cross_entropy

Adapt vocab-parallel cross entropy to the backend target signature.

API#

class core.tensor_parallel.cross_entropy.VocabParallelCrossEntropy#

Computes the Cross Entropy Loss splitting the Vocab size across tensor parallel ranks. This implementation is used in both fused and unfused cross entropy implementations

static calculate_logits_max(
vocab_parallel_logits: torch.Tensor,
) → Tuple[torch.Tensor, torch.Tensor]#

Calculates logits_max.

static gather_predicted_logits(
vocab_parallel_logits: torch.Tensor,
target: torch.Tensor,
logits_max: torch.Tensor,
vocab_start_index: int,
vocab_end_index: int,
) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor]#

Read-only part of the forward: target mask, local target ids and logits[target] - logits_max (zero for targets owned by other ranks).

Does not modify vocab_parallel_logits.

static exp_logits_inplace(
vocab_parallel_logits: torch.Tensor,
logits_max: torch.Tensor,
) → None#

In place: vocab_parallel_logits <- exp(vocab_parallel_logits - logits_max).

Pointwise only. Kept separate from any reduction over the result so that, when compiled, Inductor can keep the mutation in place (see fused_cross_entropy.py).

static calculate_predicted_logits(
vocab_parallel_logits: torch.Tensor,
target: torch.Tensor,
logits_max: torch.Tensor,
vocab_start_index: int,
vocab_end_index: int,
) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]#

Calculates predicted logits. vocab_parallel_logits is overwritten in place with exp(logits - logits_max) and returned as exp_logits.

static calculate_cross_entropy_loss(
exp_logits: torch.Tensor,
predicted_logits: torch.Tensor,
sum_exp_logits: torch.Tensor,
) → Tuple[torch.Tensor, torch.Tensor]#

Calculates cross entropy loss.

static prepare_gradient_calculation_operands(
softmax: torch.Tensor,
target_mask: torch.Tensor,
) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]#

Prepare gradient calculation operands.

static scale_gradients_inplace(
grad_input: torch.Tensor,
grad_output: torch.Tensor,
) → None#

In place: grad_input <- grad_input * grad_output.

Pointwise only. Kept separate from the target scatter so that, when compiled, Inductor can keep the mutation in place (see fused_cross_entropy.py).

static calculate_gradients(
grad_2d: torch.Tensor,
arange_1d: torch.Tensor,
masked_target_1d: torch.Tensor,
softmax_update: torch.Tensor,
grad_input: torch.Tensor,
grad_output: torch.Tensor,
) → torch.Tensor#

Calculates gradients in place in grad_input (grad_2d is a view of it).

class core.tensor_parallel.cross_entropy._VocabParallelCrossEntropy#

Bases: torch.autograd.Function

static forward(
ctx,
vocab_parallel_logits,
target,
label_smoothing=0.0,
tp_group=None,
)#

Vocab parallel cross entropy forward function.

static backward(ctx, grad_output)#

Vocab parallel cross entropy backward function.

core.tensor_parallel.cross_entropy.vocab_parallel_cross_entropy(
vocab_parallel_logits: torch.Tensor,
target: torch.Tensor,
label_smoothing: float = 0.0,
tp_group: torch.distributed.ProcessGroup | None = None,
) → torch.Tensor#

Performs cross entropy loss when logits are split across tensor parallel ranks

Parameters:

  • vocab_parallel_logits – logits split across tensor parallel ranks dimension is [sequence_length, batch_size, vocab_size/num_parallel_ranks]

  • target – correct vocab ids of dimseion [sequence_length, micro_batch_size]

  • label_smoothing – smoothing factor, must be in range [0.0, 1.0) default is no smoothing (=0.0)

  • tp_group – the tensor parallel group over which to all reduce

core.tensor_parallel.cross_entropy.unfused_cross_entropy(
logits: torch.Tensor,
labels: torch.Tensor,
tp_group: torch.distributed.ProcessGroup | None = None,
) → torch.Tensor#

Adapt vocab-parallel cross entropy to the backend target signature.

Pass the process group by keyword because the kernel’s third positional argument is label_smoothing.