|
|
from contextlib import contextmanager |
|
|
|
|
|
import torch |
|
|
import torch.distributed as dist |
|
|
from tqdm import tqdm |
|
|
|
|
|
from internlm.core.context import ParallelMode |
|
|
from internlm.core.context import global_context as gpc |
|
|
from internlm.core.scheduler import SchedulerMetricHook |
|
|
from internlm.model.metrics import AccPerplex |
|
|
|
|
|
|
|
|
@contextmanager |
|
|
def switch_evaluation_no_pipeline_scheduler(trainer, grad_accum_size, grad_accum_batch_size, metric_hook_list): |
|
|
if not gpc.is_using_pp(): |
|
|
prev_data_process_func = trainer.schedule.data_process_func |
|
|
prev_grad_accum_size = trainer.schedule._grad_accum_size |
|
|
prev_grad_accum_batch_size = trainer.schedule._grad_accum_batch_size |
|
|
prev_metric_hooks = trainer.schedule._hooks |
|
|
try: |
|
|
trainer.schedule.data_process_func = None |
|
|
trainer.schedule._grad_accum_size = grad_accum_size |
|
|
trainer.schedule._grad_accum_batch_size = grad_accum_batch_size |
|
|
trainer.schedule._hooks = metric_hook_list |
|
|
yield |
|
|
finally: |
|
|
trainer.schedule.data_process_func = prev_data_process_func |
|
|
trainer.schedule._grad_accum_size = prev_grad_accum_size |
|
|
trainer.schedule._grad_accum_batch_size = prev_grad_accum_batch_size |
|
|
trainer.schedule._hooks = prev_metric_hooks |
|
|
|
|
|
|
|
|
@contextmanager |
|
|
def switch_evaluation_pipeline_scheduler(trainer, num_microbatches, tensor_shape, metric_hook_list): |
|
|
if gpc.is_using_pp(): |
|
|
pre_data_process_func = trainer.schedule.data_process_func |
|
|
prev_num_microbatches = trainer.schedule.num_microbatches |
|
|
prev_tensor_shape = trainer.schedule.tensor_shape |
|
|
prev_metric_hooks = trainer.schedule._hooks |
|
|
try: |
|
|
trainer.schedule.data_process_func = None |
|
|
trainer.schedule.num_microbatches = num_microbatches |
|
|
trainer.schedule.tensor_shape = tensor_shape |
|
|
trainer.schedule._hooks = metric_hook_list |
|
|
yield |
|
|
finally: |
|
|
trainer.schedule.data_process_func = pre_data_process_func |
|
|
trainer.schedule.num_microbatches = prev_num_microbatches |
|
|
trainer.schedule.tensor_shape = prev_tensor_shape |
|
|
trainer.schedule._hooks = prev_metric_hooks |
|
|
|
|
|
|
|
|
@contextmanager |
|
|
def switch_sequence_parallel_mode(): |
|
|
prev_mode = gpc.config.parallel.sequence_parallel |
|
|
try: |
|
|
gpc.config.parallel.sequence_parallel = False |
|
|
yield |
|
|
finally: |
|
|
gpc.config.parallel.sequence_parallel = prev_mode |
|
|
|
|
|
|
|
|
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 not streaming and len(val_dl) == 0 and verbose: |
|
|
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 |
|
|
grad_accum_size = total_val_bsz // data_cfg.micro_bsz |
|
|
grad_accum_batch_size = data_cfg.micro_bsz |
|
|
with switch_evaluation_no_pipeline_scheduler( |
|
|
trainer=trainer, |
|
|
grad_accum_size=grad_accum_size, |
|
|
grad_accum_batch_size=grad_accum_batch_size, |
|
|
metric_hook_list=[val_sche_metric_hook], |
|
|
): |
|
|
_, _, loss = trainer.execute_schedule( |
|
|
batch, forward_only=True, return_loss=True, return_output_label=False |
|
|
) |
|
|
if verbose: |
|
|
if isinstance(loss, dict): |
|
|
loss = sum(loss.values()) |
|
|
val_loss += loss.item() |
|
|
|
|
|
assert val_idx != -1 |
|
|
dist.barrier() |
|
|
|
|
|
val_res = val_metric.get_metric() |
|
|
if verbose and (streaming or len(val_dl) != 0): |
|
|
val_loss = val_loss / (val_idx + 1 + 1e-6) |
|
|
infos = { |
|
|
"step": step_count, |
|
|
f"val/{val_name}_loss": val_loss, |
|
|
f"val/{val_name}_acc": val_res["acc"], |
|
|
f"val/{val_name}_plex": val_res["perplexity"], |
|
|
} |
|
|
|
|
|
for key, value in infos.items(): |
|
|
writer.add_scalar(key=key, value=value, step=step_count) |
|
|
|
|
|
if update_panel: |
|
|
logger.info( |
|
|
f"Validation on {val_name}: " + " ".join([f"{key}={value}" for key, value in infos.items()]), |
|
|
extra={ |
|
|
"step": step_count, |
|
|
"val_loss": val_loss, |
|
|
"val_acc": val_res["acc"], |
|
|
"val_perplexity": val_res["perplexity"], |
|
|
}, |
|
|
) |
|
|
else: |
|
|
logger.info( |
|
|
f"Validation on {val_name}: " + " ".join([f"{key}={value}" for key, value in infos.items()]) |
|
|
) |
|
|
|
|
|
trainer.train() |
|
|
torch.cuda.empty_cache() |
|
|
dist.barrier() |
|
|
|
