Function evaluate

codegeex/megatron/training.py:1022–1087

Evaluation.

(forward_step_func, data_iterator, model, verbose=False)

def evaluate(forward_step_func, data_iterator, model, verbose=False):
    """Evaluation."""
    args = get_args()

    # Turn on evaluation mode which disables dropout.
    for model_module in model:
        model_module.eval()

    total_loss_dict = {}

    with torch.no_grad():
        iteration = 0
        while iteration < args.eval_iters:
            iteration += 1
            if verbose and iteration % args.log_interval == 0:
                print_rank_0("Evaluating iter {}/{}".format(iteration, args.eval_iters))

            if mpu.get_pipeline_model_parallel_world_size() > 1:
                if args.virtual_pipeline_model_parallel_size is not None:
                    forward_backward_func = (
                        forward_backward_pipelining_with_interleaving
                    )
                else:
                    forward_backward_func = (
                        forward_backward_pipelining_without_interleaving
                    )
            else:
                forward_backward_func = forward_backward_no_pipelining

            if args.deepspeed and not args.no_pipeline_parallel:
                # DeepSpeed uses eval_batch() and already aggregates losses.
                assert isinstance(model, list) and len(model) == 1
                loss = model[0].eval_batch(data_iterator)
                loss_dicts = [{"lm loss": loss}] * get_num_microbatches()
            else:
                loss_dicts = forward_backward_func(
                    forward_step_func,
                    data_iterator,
                    model,
                    optimizer=None,
                    timers=None,
                    forward_only=True,
                )

            if mpu.is_pipeline_last_stage(ignore_virtual=True):
                # Reduce across processes.
                for loss_dict in loss_dicts:
                    for key in loss_dict:
                        total_loss_dict[key] = (
                            total_loss_dict.get(key, torch.cuda.FloatTensor([0.0]))
                            + loss_dict[key]
                        )

            args.consumed_valid_samples += (
                mpu.get_data_parallel_world_size()
                * args.micro_batch_size
                * get_num_microbatches()
            )