File size: 2,140 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
import torch
import deepspeed.comm as dist
class _VocabSequenceParallelCrossEntropy(torch.autograd.Function):
@staticmethod
def forward(ctx, vocab_seq_parallel_logits, target, sp_group):
# vocab_seq_parallel_logits: [S/P, B, V]
# target: [S/P, B]
# return: [S, B]
# Need softmax for backward
softmax = torch.nn.functional.softmax(vocab_seq_parallel_logits, dim=-1)
ctx.vocab_size = vocab_seq_parallel_logits.size(2)
loss = torch.nn.functional.nll_loss(softmax.log().view(-1, ctx.vocab_size), target.view(-1), reduction='none')
sp_world_size = dist.get_world_size(sp_group)
sp_rank = dist.get_rank(sp_group)
ctx.sp_world_size = sp_world_size
ctx.sp_rank = sp_rank
ctx.seqlen = vocab_seq_parallel_logits.size(0) * sp_world_size
batch_size = vocab_seq_parallel_logits.size(1)
loss_all = torch.empty(ctx.seqlen,
batch_size,
dtype=vocab_seq_parallel_logits.dtype,
device=vocab_seq_parallel_logits.device)
dist.all_gather_into_tensor(loss_all, loss, group=sp_group)
ctx.save_for_backward(softmax, target)
return loss_all
@staticmethod
def backward(ctx, grad_output):
softmax, target = ctx.saved_tensors
step_seqlen = ctx.seqlen // ctx.sp_world_size
sp_rank = ctx.sp_rank
grad_output_part = grad_output[step_seqlen * sp_rank:step_seqlen * (sp_rank + 1), :]
grad_input = softmax
grad_2d = grad_input.view(-1, ctx.vocab_size)
arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device)
grad_2d[arange_1d, target.view(-1)] -= 1
grad_input.mul_(grad_output_part.unsqueeze(dim=-1))
return grad_input, None, None, None
def vocab_sequence_parallel_cross_entropy(vocab_parallel_logits, target, sp_group):
return _VocabSequenceParallelCrossEntropy.apply(vocab_parallel_logits, target, sp_group)
|