mirror of https://github.com/InternLM/InternLM
344 lines
12 KiB
Python
344 lines
12 KiB
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
import socket
|
|
import time
|
|
import traceback
|
|
from functools import partial
|
|
from typing import List, Optional
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
import internlm
|
|
from internlm.core.context import ParallelMode
|
|
from internlm.core.context import global_context as gpc
|
|
from internlm.core.scheduler import SchedulerHook
|
|
from internlm.core.trainer import TrainState
|
|
from internlm.initialize import initialize_distributed_env
|
|
from internlm.model.loss import FlashGPTLMLoss
|
|
from internlm.model.metrics import AccPerplex, SchedulerMetricHook
|
|
from internlm.model.overlap_handler import FSTPOverlapSchedulerHook
|
|
from internlm.monitor import initialize_monitor_manager, send_alert_message
|
|
from internlm.monitor.monitor import monitor_manager as mm
|
|
from internlm.train import (
|
|
get_train_data_loader,
|
|
get_validation_data_loader,
|
|
initialize_llm_profile,
|
|
initialize_model,
|
|
initialize_optimizer,
|
|
load_new_batch,
|
|
record_current_batch_training_metrics,
|
|
)
|
|
from internlm.utils.common import (
|
|
BatchSkipper,
|
|
get_megatron_flops,
|
|
launch_time,
|
|
parse_args,
|
|
)
|
|
from internlm.utils.evaluation import evaluate_on_val_dls
|
|
from internlm.utils.gputest import empty_cache_and_diag
|
|
from internlm.utils.logger import get_logger, initialize_uniscale_logger
|
|
from internlm.utils.megatron_timers import megatron_timer as timer
|
|
from internlm.utils.model_checkpoint import CheckpointManager
|
|
from internlm.utils.parallel import get_parallel_log_file_name
|
|
from internlm.utils.simple_memory_profiler import SimpleMemoryProfiler
|
|
from internlm.utils.writer import Writer
|
|
|
|
# global llm logger
|
|
logger = get_logger(__file__)
|
|
|
|
|
|
def initialize_llm_logger(start_time: str):
|
|
"""
|
|
Initialize customed uniscale logger.
|
|
|
|
Args:
|
|
start_time (str): The launch time of current training job.
|
|
|
|
Returns: The instance of uniscale logger.
|
|
"""
|
|
|
|
uniscale_logger = initialize_uniscale_logger(
|
|
job_name=gpc.config.JOB_NAME, launch_time=start_time, file_name=get_parallel_log_file_name()
|
|
)
|
|
if uniscale_logger is not None:
|
|
global logger
|
|
logger = uniscale_logger
|
|
|
|
return uniscale_logger
|
|
|
|
|
|
def get_scheduler_hooks(
|
|
metric: Optional[AccPerplex] = None, activation_checkpoint: bool = False
|
|
) -> List[SchedulerHook]:
|
|
scheduler_hooks: List[SchedulerHook] = []
|
|
|
|
if metric is not None:
|
|
scheduler_hooks.append(
|
|
SchedulerMetricHook(
|
|
metric=metric,
|
|
skip=(
|
|
gpc.is_using_pp()
|
|
and hasattr(gpc.config.model, "num_chunks")
|
|
and gpc.config.model.num_chunks > 1
|
|
and gpc.config.parallel["pipeline"].get("interleaved_overlap", False)
|
|
),
|
|
),
|
|
)
|
|
|
|
if activation_checkpoint:
|
|
scheduler_hooks.append(FSTPOverlapSchedulerHook(gpc.fstp_handler))
|
|
|
|
return scheduler_hooks
|
|
|
|
|
|
def main(args):
|
|
# init setting
|
|
skip_batches = gpc.config.data.skip_batches
|
|
total_steps = gpc.config.data.total_steps
|
|
valid_every = gpc.config.data.valid_every
|
|
label_smoothing = gpc.config.loss.label_smoothing
|
|
|
|
get_tflops_func = partial(
|
|
get_megatron_flops,
|
|
checkpoint=gpc.config.model.checkpoint,
|
|
seq_len=gpc.config.SEQ_LEN,
|
|
hidden_size=gpc.config.model.hidden_size,
|
|
num_layers=gpc.config.model.num_layers,
|
|
vocab_size=gpc.config.model.vocab_size,
|
|
global_batch_size=gpc.config.data.micro_bsz * gpc.config.data.micro_num * gpc.get_world_size(ParallelMode.DATA),
|
|
global_world_size=gpc.get_world_size(ParallelMode.GLOBAL),
|
|
mlp_ratio=gpc.config.MLP_RATIO,
|
|
)
|
|
|
|
# get and broadcast current time
|
|
current_time = launch_time()
|
|
objs = [current_time]
|
|
dist.broadcast_object_list(objs, src=0)
|
|
current_time = objs[0]
|
|
|
|
# initialize customed llm logger
|
|
uniscale_logger = initialize_llm_logger(start_time=current_time)
|
|
|
|
# initialize model
|
|
model = initialize_model()
|
|
|
|
with open(args.config, "r") as f:
|
|
config_lines = f.readlines()
|
|
|
|
# initialize loss function
|
|
criterion = FlashGPTLMLoss(parallel_output=True, label_smoothing=label_smoothing)
|
|
|
|
# initialize the train and validation data loader
|
|
train_dl, dataset_types = get_train_data_loader(num_worker=4)
|
|
val_dls = get_validation_data_loader()
|
|
|
|
# initialize and resume train state
|
|
train_state = TrainState(gpc.config, train_dl.batch_sampler)
|
|
|
|
optimizer, beta2_scheduler, lr_scheduler = initialize_optimizer(model=model)
|
|
|
|
ckpt_manager = CheckpointManager(
|
|
ckpt_config=gpc.config.ckpt,
|
|
model=model,
|
|
optimizer=optimizer,
|
|
lr_scheduler=lr_scheduler,
|
|
train_dl=train_dl,
|
|
model_config=gpc.config.model,
|
|
model_config_file="".join(config_lines),
|
|
feishu_address=gpc.config.monitor.alert.feishu_alert_address,
|
|
)
|
|
|
|
# Loading other persistent training states.
|
|
ckpt_manager.try_resume_training(train_state, current_time)
|
|
|
|
# initialize customed llm writer
|
|
writer = Writer(
|
|
job_name=gpc.config.JOB_NAME,
|
|
launch_time=current_time,
|
|
file_name=get_parallel_log_file_name(),
|
|
tensorboard_folder=gpc.config.tensorboard_folder,
|
|
resume_tb_folder=train_state.resume_tb_folder, # resume from ckpt.
|
|
step_count=train_state.step_count, # resume from ckpt.
|
|
config=config_lines,
|
|
logger=logger,
|
|
enable_tb=gpc.config.enable_tb,
|
|
)
|
|
|
|
# initialize metric for calculating accuracy and perplexity
|
|
metric = AccPerplex(
|
|
device=torch.cuda.current_device(),
|
|
tp_pg=gpc.get_group(ParallelMode.TENSOR),
|
|
dp_pg=gpc.get_group(ParallelMode.DATA),
|
|
dataset_types=dataset_types,
|
|
)
|
|
|
|
# initialize trainer
|
|
|
|
trainer, train_dl, _, _ = internlm.initialize_trainer(
|
|
model=model,
|
|
optimizer=optimizer,
|
|
criterion=criterion,
|
|
train_dataloader=train_dl,
|
|
lr_scheduler=lr_scheduler,
|
|
beta2_scheduler=beta2_scheduler,
|
|
scheduler_hooks=get_scheduler_hooks(metric, gpc.config.model.checkpoint),
|
|
)
|
|
|
|
# initialize simple memory profiler
|
|
if args.profiling:
|
|
memory_profiler = SimpleMemoryProfiler(
|
|
model,
|
|
optimizer.optim,
|
|
log_folder=f"memory_trace/rank{gpc.get_global_rank()}_"
|
|
+ f"dp{gpc.get_local_rank(ParallelMode.DATA)}_"
|
|
+ f"tp{gpc.get_local_rank(ParallelMode.TENSOR)}",
|
|
)
|
|
else:
|
|
memory_profiler = None
|
|
|
|
# initialize the batch skipper
|
|
batch_skipper = BatchSkipper(skip_batches)
|
|
|
|
trainer.train()
|
|
|
|
# transfer the train data loader into train data iterator
|
|
train_iter = iter(train_dl)
|
|
|
|
with initialize_llm_profile(profiling=args.profiling, start_time=current_time) as prof:
|
|
# start iterating the train data and begin training
|
|
for batch_count in range(train_state.batch_count, total_steps):
|
|
empty_cache_and_diag(batch_count, interval=gpc.config.data.empty_cache_and_diag_interval)
|
|
# torch.cuda.memory._record_memory_history()
|
|
start_time = time.time()
|
|
timer("one-batch").start()
|
|
|
|
# load batch data
|
|
batch, train_iter = load_new_batch(train_dl=train_dl, train_iter=train_iter, train_state=train_state)
|
|
|
|
# record the consumed samples in training
|
|
train_state.batch_count = batch_count
|
|
train_state.num_consumed_samples_in_epoch += len(batch[1])
|
|
if batch_skipper(batch_count): # skip this batch
|
|
if gpc.is_rank_for_log():
|
|
logger.info(f"Skip batch count:`{batch_count}`...")
|
|
timer("one-batch").stop()
|
|
continue
|
|
|
|
# zero the grads of parameters
|
|
trainer.zero_grad()
|
|
# process data
|
|
if batch[0].get("type_ids", None) is not None:
|
|
metric.set_current_type_ids(type_ids=batch[0].pop("type_ids", None))
|
|
|
|
# do forward and backward
|
|
timer("fwd-bwd").start()
|
|
|
|
moe_loss = None
|
|
if hasattr(gpc.config.model, "num_experts"):
|
|
_, _, loss, moe_loss = trainer.execute_schedule(
|
|
batch,
|
|
forward_only=False,
|
|
return_loss=True,
|
|
return_output_label=False,
|
|
)
|
|
else:
|
|
_, _, loss = trainer.execute_schedule(
|
|
batch,
|
|
forward_only=False,
|
|
return_loss=True,
|
|
return_output_label=False,
|
|
)
|
|
timer("fwd-bwd").stop()
|
|
|
|
# update parameters, and returns (success_update, grad_norm)
|
|
trainer_result = trainer.step()
|
|
assert trainer_result is not None
|
|
|
|
success_update, grad_norm_groups = trainer_result
|
|
if success_update: # update parameters successfully
|
|
train_state.step_count += 1
|
|
else:
|
|
train_state.inf_nan_skip_batches += 1 # record the amount of updating parameters unsuccessfully.
|
|
if -1 in grad_norm_groups.values() and gpc.is_rank_for_log(): # -1 encodes a specific failure case
|
|
logger.warning(f"Warning: skip parameter update at step {batch_count}.")
|
|
send_alert_message(
|
|
address=gpc.config.monitor.alert.feishu_alert_address,
|
|
message=f"Warning: skip parameter update at step {batch_count}.",
|
|
)
|
|
|
|
# calculate and record the training metrics, eg. loss, accuracy and so on.
|
|
record_current_batch_training_metrics(
|
|
get_tflops_func=get_tflops_func,
|
|
logger=logger,
|
|
writer=writer,
|
|
success_update=success_update,
|
|
batch_count=batch_count,
|
|
batch=batch,
|
|
train_state=train_state,
|
|
optimizer=optimizer,
|
|
beta2_scheduler=beta2_scheduler,
|
|
trainer=trainer,
|
|
start_time=start_time,
|
|
loss=loss,
|
|
moe_loss=moe_loss,
|
|
grad_norm=grad_norm_groups,
|
|
metric=metric,
|
|
update_panel=uniscale_logger is not None,
|
|
)
|
|
|
|
timer("one-batch").stop()
|
|
|
|
# evaluate on validation data loaders
|
|
if valid_every > 0 and train_state.step_count % valid_every == 0:
|
|
evaluate_on_val_dls(
|
|
trainer=trainer,
|
|
val_dls=val_dls,
|
|
writer=writer,
|
|
logger=logger,
|
|
step_count=train_state.step_count,
|
|
update_panel=uniscale_logger is not None,
|
|
)
|
|
|
|
# checkpoint the training states in specific steps, which is determined by the args "checkpoint_every"
|
|
# # save batch sampler that tracks the true consumed samples
|
|
now_break = ckpt_manager.try_save_checkpoint(train_state)
|
|
if now_break:
|
|
break
|
|
|
|
if memory_profiler is not None:
|
|
memory_profiler.step()
|
|
|
|
prof.step()
|
|
|
|
if gpc.fstp_handler is not None:
|
|
gpc.fstp_handler.clear_memory_pool()
|
|
# torch.cuda.memory._dump_snapshot(f"my_snapshot_{gpc.get_global_rank()}.pickle")
|
|
torch.cuda.reset_peak_memory_stats()
|
|
|
|
ckpt_manager.wait_async_upload_finish()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
hostname = socket.gethostname()
|
|
|
|
# initialize distributed environment
|
|
initialize_distributed_env(config=args.config, launcher=args.launcher, master_port=args.port, seed=args.seed)
|
|
assert hasattr(gpc, "config") and gpc.config is not None
|
|
|
|
# initialize monitor manager context
|
|
with initialize_monitor_manager(
|
|
job_name=gpc.config.JOB_NAME, alert_address=gpc.config.monitor.alert.feishu_alert_address
|
|
):
|
|
try:
|
|
main(args)
|
|
except Exception:
|
|
logger.error(
|
|
f"Raise exception from {hostname} with rank id: {gpc.get_global_rank()}\n{traceback.format_exc()}",
|
|
)
|
|
mm.monitor_exception(
|
|
alert_address=gpc.config.monitor.alert.feishu_alert_address, excp_info=traceback.format_exc()
|
|
)
|