Function train_step

Megatron-LM/pretrain_gpt2.py:357–388

Single training step.

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

        #input("Press Any Key To Continue ..")

def train_step(data_iterator, model, optimizer, lr_scheduler,
               args, timers):
    """Single training step."""

    # Forward model for one step.
    timers('forward').start()
    lm_loss = forward_step(data_iterator, model, args, timers)
    timers('forward').stop()

    #print_rank_0("loss is {}".format(lm_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 = 0
    timers('optimizer').start()
    if args.deepspeed:
        model.step()
    else:
        optimizer.step()

        # Update learning rate.
        if not (args.fp16 and optimizer.overflow):
            lr_scheduler.step()
        else:
            skipped_iter = 1
    timers('optimizer').stop()

    return lm_loss_reduced, skipped_iter


def train(model, optimizer, lr_scheduler,