Repository: hpcaitech/ColossalAI
File: colossalai/shardformer/layer/loss.py (Python)

Code before:
        masked_target[mask] = 0
        masked_target_1d = masked_target.view(-1).contiguous()
        # wait for allreduce
        handle.wait()
        vocab_logits = vocab_logits - logits_max.unsqueeze(dim=-1)
        # cal exp_logits
        exp_logits = torch.exp(vocab_logits)
        # cal local sum_exp_logits
        sum_exp_logits = torch.sum(exp_logits, dim=-1, dtype=torch.float32)
        # all_reduce get global sum_exp_logits
        dist.all_reduce(sum_exp_logits, op=dist.ReduceOp.SUM, group=process_group)
        # cal log_softmax
        log_probs = vocab_logits - torch.log(sum_exp_logits.unsqueeze(dim=-1))

        ######
        # 2.gather via labels
        ######
        log_probs = log_probs.gather(dim=-1, index=masked_target.unsqueeze(-1))
        print(f"log_probs in ops {log_probs.shape}")
        # set masked val to zero, then all reduce
        log_probs[mask.unsqueeze(-1)] = 0
        # allreduce log_probs with ops SUM
        dist.all_reduce(log_probs, op=dist.ReduceOp.SUM, group=process_group)
        ctx.save_for_backward(exp_logits, mask, masked_target_1d, sum_exp_logits)
        ctx.dtype = dtype
        return log_probs

    @staticmethod
    def backward(ctx, grad_output):
        exp_logits, mask, masked_target_1d, sum_exp_logits = ctx.saved_tensors
        softmax_logits = exp_logits / sum_exp_logits.unsqueeze(dim=-1)
        partion_vocab_size = softmax_logits.shape[-1]
        softmax_logits_2d = softmax_logits.view(-1, partion_vocab_size)
        update = 1.0 - mask.view(-1).float().to(ctx.dtype)
        softmax_logits_2d[torch.arange(0, softmax_logits_2d.shape[0]), masked_target_1d] -= update
        grad_logits = -softmax_logits.mul_(grad_output)
        return grad_logits, None, None, None, None, None, None


def cross_entropy_1d(
    vocab_logits: torch.Tensor,
    labels: torch.Tensor,
    ignore_index: int = _IGNORE_IDX,
    process_group: ProcessGroup = None,
    vocab_size: int = None,
    dtype: torch.dtype = None,
    mode: str = "mean",
) -> torch.Tensor:
    return DistCrossEntropy.apply(vocab_logits, labels, ignore_index, process_group, vocab_size, dtype, mode)



Review comment (suggestion):
Please remove unnecessary comments. For example, `allreduce log_probs with ops SUM` or  `cal exp_logits`, it is not necessary to explain this simple sentence. You should explain you algorithm (how to do forward here). 



Code after:
            up_threshold = global_vocab_size
        # mask
        mask = (target < down_threshold) | (target >= up_threshold)
        masked_target = target.clone() - down_threshold
        masked_target[mask] = 0
        masked_target_1d = masked_target.view(-1).contiguous()
        handle.wait()

        ##################
        # Step3:Calculate global summation exp logits
        ##################
        vocab_logits = vocab_logits - logits_max.unsqueeze(dim=-1)
        exp_logits = torch.exp(vocab_logits)
        sum_exp_logits = torch.sum(exp_logits, dim=-1, dtype=torch.float32)  # local summation exp logits
        dist.all_reduce(sum_exp_logits, op=dist.ReduceOp.SUM, group=process_group)

        ##################
        # Step4:Calculate local prob. We first cal log_softmax, then select log probs via local mask
        ##################
        log_probs = vocab_logits - torch.log(sum_exp_logits.unsqueeze(dim=-1))  # cal log_softmax
        log_probs = log_probs.gather(dim=-1, index=masked_target.unsqueeze(-1))
        log_probs[mask.unsqueeze(-1)] = 0  # set masked val to zero
        dist.all_reduce(log_probs, op=dist.ReduceOp.SUM, group=process_group)

        ctx.save_for_backward(exp_logits, mask, masked_target_1d, sum_exp_logits)
        ctx.dtype = dtype
        return log_probs

    @staticmethod
    def backward(ctx, grad_output):
        exp_logits, mask, masked_target_1d, sum_exp_logits = ctx.saved_tensors
        ##################
        # Step1:Find the global sofmax value
        ##################
        softmax_logits = exp_logits / sum_exp_logits.unsqueeze(dim=-1)

        ##################
        # Step2:Update softmax value based on local target index
        ##################
        partion_vocab_size = softmax_logits.shape[-1]
        softmax_logits_2d = softmax_logits.view(-1, partion_vocab_size)
        update = 1.0 - mask.view(-1).float().to(ctx.dtype)
        softmax_logits_2d[torch.arange(0, softmax_logits_2d.shape[0]), masked_target_1d] -= update

        ##################
        # Step3:Calculate grad_output, which is the gradient of the loss function with respect to the output of logsoftmax
        ##################
        grad_logits = -softmax_logits.mul_(grad_output)
        return grad_logits, None, None, None, None, None, None