mirror of https://github.com/hpcaitech/ColossalAI
42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
from colossalai.core import global_context as gpc
|
||
|
|
from colossalai.context import ParallelMode
|
||
|
|
from colossalai.logging import get_dist_logger
|
||
|
|
import torch.nn.functional as F
|
||
|
|
import torch.distributed as dist
|
||
|
|
from .cross_entropy import vocab_cross_entropy
|
||
|
|
|
||
|
|
|
||
|
|
class BertLoss(nn.Module):
|
||
|
|
|
||
|
|
def forward(self,
|
||
|
|
lm_loss,
|
||
|
|
sop_logits,
|
||
|
|
loss_mask,
|
||
|
|
sentence_order):
|
||
|
|
lm_loss_ = lm_loss.float()
|
||
|
|
loss_mask = loss_mask.float()
|
||
|
|
loss_mask_sum = loss_mask.sum()
|
||
|
|
lm_loss = torch.sum(
|
||
|
|
lm_loss_.view(-1) * loss_mask.reshape(-1))
|
||
|
|
|
||
|
|
lm_loss /= loss_mask_sum
|
||
|
|
|
||
|
|
torch.distributed.all_reduce(
|
||
|
|
lm_loss,
|
||
|
|
group=gpc.get_group(ParallelMode.SEQUENCE)
|
||
|
|
)
|
||
|
|
|
||
|
|
if sop_logits is not None:
|
||
|
|
sop_loss = F.cross_entropy(sop_logits.view(-1, 2).float(),
|
||
|
|
sentence_order.view(-1),
|
||
|
|
ignore_index=-1)
|
||
|
|
sop_loss = sop_loss.float()
|
||
|
|
loss = lm_loss + sop_loss * gpc.get_world_size(ParallelMode.SEQUENCE)
|
||
|
|
else:
|
||
|
|
sop_loss = None
|
||
|
|
loss = lm_loss
|
||
|
|
|
||
|
|
return loss
|