Function train_step

pretrain_gpt2.py:406–448

Single training step.

(data_iterator, model, optimizer, lr_scheduler,
               args, timers, mems)

        print(" ")

def train_step(data_iterator, model, optimizer, lr_scheduler,
               args, timers, mems):
    """Single training step."""
    while True:
        # Forward model for one step.
        timers('forward').start()
        lm_loss, mems, img_loss, txt_loss = forward_step(data_iterator, model, args, timers, mems)
        timers('forward').stop()

        if (img_loss + txt_loss).isnan().any() or (img_loss + txt_loss).isinf().any():
            print('Skipping backward and optimizer step for nan or inf in forwarding!')
            return (img_loss + txt_loss), 1, mems, img_loss, txt_loss

        # Calculate gradients, reduce across processes, and clip.
        timers('backward').start()
        lm_loss_reduced = backward_step(optimizer, model, lm_loss, args, timers)
        timers('backward').stop()

        # Update parameters.
        skipped_iter, complete = 0, False
        timers('optimizer').start()
        if args.deepspeed:
            if model.is_gradient_accumulation_boundary():
                model.step()
                complete = True
                if not (args.fp16 and optimizer.overflow):
                    lr_scheduler.step()
                else:
                    skipped_iter = 1
            else:
                model.step()
        else:
            optimizer.step()
            complete = True
            # Update learning rate.
            if not (args.fp16 and optimizer.overflow):
                lr_scheduler.step()
            else:
                skipped_iter = 1
        timers('optimizer').stop()
        if complete:
            break
    return lm_loss_reduced, skipped_iter, mems, img_loss, txt_loss


def report_iteration_metrics(summary_writer, optimizer, lr, loss, elapsed_time, step, total_step, args, img_loss, txt_loss):