Function train

codegeex/megatron/training.py:882–1019

Train the model function.

(
    forward_step_func,
    valid_forward_step_func,
    model,
    optimizer,
    lr_scheduler,
    train_data_iterator,
    valid_data_iterator,
)

def train(
    forward_step_func,
    valid_forward_step_func,
    model,
    optimizer,
    lr_scheduler,
    train_data_iterator,
    valid_data_iterator,
):
    """Train the model function."""
    args = get_args()
    timers = get_timers()

    # Write args to tensorboard
    write_args_to_tensorboard()

    if args.wandb_logging:
        torch.distributed.barrier()
        print_datetime("before the initialization of wandb")
        timers("wandb-init").start()
        if is_last_rank():
            initialize_wandb_experiment()
        torch.distributed.barrier()
        timers("wandb-init").stop()
        timers.log(["wandb-init"])

    # Turn on training mode which enables dropout.
    for model_module in model:
        model_module.train()

    # Tracking loss.
    total_loss_dict = {}

    # Iterations.
    iteration = args.iteration

    timers("interval-time").start()
    print_datetime("before the start of training step")
    report_memory_flag = True
    
    while iteration < args.train_iters and (
        args.train_tokens is None or args.consumed_train_tokens < args.train_tokens
    ):
        # print_rank_0(f'=> iteration {iteration}')
        update_num_microbatches(args.consumed_train_samples)
        if args.deepspeed:
            # inform deepspeed of any batch size changes
            global_batch_size = (
                mpu.get_data_parallel_world_size()
                * args.micro_batch_size
                * get_num_microbatches()
            )
            model[0].set_train_batch_size(global_batch_size)

        # print_rank_0(f"==> running train step for iteration {iteration}")
        loss_dict, skipped_iter, grad_norm, num_zeros_in_grad = train_step(
            forward_step_func, train_data_iterator, model, optimizer, lr_scheduler
        )