mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
42 lines
1.2 KiB
42 lines
1.2 KiB
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
|