Function train_step

DeepSpeedExample/megatron/training.py:361–395

Single training step.

(forward_step_func, data_iterator,
               model, optimizer, lr_scheduler)

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

    #see_memory_usage(f'before forward {model.global_steps}', force=True)
    # Forward model for one step.
    timers('forward').start()
    loss, loss_reduced = forward_step_func(data_iterator, model)
    timers('forward').stop()

    #see_memory_usage(f'before backward {model.global_steps}', force=True)
    # Calculate gradients, reduce across processes, and clip.
    timers('backward').start()
    backward_step(optimizer, model, loss)
    timers('backward').stop()

    #print('[Rank=',torch.distributed.get_rank(),'] backward complete!')
    #see_memory_usage(f'before optimizer {model.global_steps}', force=True)
    # 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 loss_reduced, skipped_iter


def training_log(loss_dict, total_loss_dict, learning_rate, iteration,