Function evaluate_on_val_dls

internlm/utils/evaluation.py:63–168

(
    trainer,
    val_dls,
    writer,
    logger,
    step_count,
    update_panel: bool = False,
    streaming: bool = False,
)

def evaluate_on_val_dls(
    trainer,
    val_dls,
    writer,
    logger,
    step_count,
    update_panel: bool = False,
    streaming: bool = False,
):
    with switch_sequence_parallel_mode():
        torch.cuda.empty_cache()
        trainer.eval()
        verbose = gpc.is_rank_for_log()
        data_cfg = gpc.config.data

        for val_name, val_dl in val_dls.items():
            if len(val_dl) == 0 and verbose and not streaming:
                logger.info(f"Validation dataset: {val_name} is empty")
                continue

            val_metric = AccPerplex(
                device=torch.cuda.current_device(),
                tp_pg=gpc.get_group(ParallelMode.TENSOR),
                dp_pg=gpc.get_group(ParallelMode.DATA),
            )
            val_sche_metric_hook = SchedulerMetricHook(metric=val_metric)

            val_loss = 0
            val_idx = -1
            for val_idx, batch in tqdm(
                enumerate(val_dl),
                desc="Val.",
                total=len(val_dl) if not streaming else None,
                position=1,
                disable=not verbose,
                leave=False,
            ):
                with torch.inference_mode():
                    if gpc.is_using_pp():
                        total_val_bsz = len(batch[1])
                        assert total_val_bsz % data_cfg.micro_bsz == 0
                        num_microbatches = total_val_bsz // data_cfg.micro_bsz
                        tensor_shape = torch.Size(
                            [data_cfg.micro_bsz, batch[0]["input_ids"].shape[1], gpc.config.HIDDEN_SIZE]
                        )

                        with switch_evaluation_pipeline_scheduler(
                            trainer=trainer,
                            num_microbatches=num_microbatches,
                            tensor_shape=tensor_shape,
                            metric_hook_list=[val_sche_metric_hook],
                        ):
                            _, _, loss = trainer.execute_schedule(
                                batch, forward_only=True, return_loss=True, return_output_label=False
                            )
                    else:
                        total_val_bsz = len(batch[1])
                        assert total_val_bsz % data_cfg.micro_bsz == 0