I am trying to train a DDP model (one GPU per process, but I’ve added the with autocast(enabled=args.use_mp): to model forward just in case) with mixed precision using torch.cuda.amp with train_bert function.
The model trains fine without amp as well as with autocast(enabled=False). When I try running it with mixed precision (args.use_mp = True), I get nan loss after first iteration.
I used autograd.detect_anomaly() to find that nan occurs in CrossEntropyLoss: RuntimeError: Function ‘LogSoftmaxBackward’ returned nan values in its 0th output. Not sure what kind of mistake am I looking for.
Below I insert the code for the training function and to the LossClass (criterion):
def train_bert(rank, args, epoch,
model, optimizer, criterion, loss_components,
train_generator, phase, scaler=None):
"""
bert training on args.bert_steps batches
:param rank: process rank
:param args: config class
:param epoch: epoch number for logging
:param model: bert model
:param optimizer: bert optimizer
:param criterion: callable loss
:param loss_components: torch.Tensor with cumulative losses for each objective
:param train_generator: train data loader (chunk wise)
:param phase: phase for phase wise training
:param scaler: scaler for mixed precision training (used if args.use_apex=True)
:return: mean step loss on the part of the training set
"""
model.train()
optimizer.zero_grad()
nb_tr_steps = phase * args.bert_steps
time_dict = {'step_time': 0.0, 'forward_time': 0.0, 'backward_time': 0.0}
for step, batch in islice(enumerate(tqdm(train_generator)), phase * args.bert_steps, (phase + 1) * args.bert_steps):
inp, lbl, meta = batch
X_ids, X_type, X_attn = (X.cuda(rank) for X in inp)
lbl = lbl.cuda(rank, non_blocking=True)
start = time.perf_counter()
with autocast(enabled=args.use_mp):
pred = model(X_ids, X_type, X_attn)
time_dict['forward_time'] += time.perf_counter() - start
components, loss = criterion(pred, lbl, args)
loss = loss / args.grad_acum_steps
start = time.perf_counter()
if args.use_mp:
scaler.scale(loss).backward()
else:
loss.backward()
time_dict['backward_time'] += time.perf_counter() - start
loss_components += components
nb_tr_steps += 1
if nb_tr_steps % args.grad_acum_steps == 0:
if args.use_mp:
scaler.step(optimizer)
scaler.update()
else:
torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip_grad)
optimizer.step()
optimizer.zero_grad()
return loss_components.sum() / nb_tr_steps
class ConditionalLoss:
def __init__(self, args, rank):
self.bce = nn.BCEWithLogitsLoss().cuda(rank)
self.ce = nn.CrossEntropyLoss(reduction='none').cuda(rank)
self.loss_weights = args.loss_weights
def _binary_loss(self, pred, lbl):
return self.bce(pred, lbl.unsqueeze(1))
def _start_end_loss(self, pred, lbl, is_yes_no):
is_span = (1 - is_yes_no)
return (is_span * (self.ce(pred[0], lbl[0]) + self.ce(pred[1], lbl[1]))).mean()
def __call__(self, preds, labels, args):
pred_is_yes_no, pred_span = preds
is_yes_no, yes_no, span = unpack_lbls(labels)
s_w, yn_w, iyn_w = self.loss_weights
is_y_n_loss = iyn_w * self._binary_loss(pred_is_yes_no, is_yes_no)
span_loss = s_w * self._start_end_loss(pred_span, span, is_yes_no)
return torch.Tensor([span_loss.item(), is_y_n_loss.item()]), is_y_n_loss + span_loss