紫苏
GitHub
parent
ab170f57f3
commit
9df92b986c
|
|
@ -30,10 +30,12 @@ def get_lr(current_step, total_steps, lr):
|
|||
return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
|
||||
|
||||
|
||||
def train_epoch(epoch, wandb):
|
||||
def train_epoch(epoch, wandb, start_step=0):
|
||||
loss_fct = nn.CrossEntropyLoss(reduction='none')
|
||||
start_time = time.time()
|
||||
for step, (X, Y, loss_mask) in enumerate(train_loader):
|
||||
if step < start_step:
|
||||
continue
|
||||
X = X.to(args.device)
|
||||
Y = Y.to(args.device)
|
||||
loss_mask = loss_mask.to(args.device)
|
||||
|
|
@ -90,6 +92,15 @@ def train_epoch(epoch, wandb):
|
|||
state_dict = model.state_dict()
|
||||
state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存
|
||||
torch.save(state_dict, ckp)
|
||||
|
||||
# 保存优化器状态、当前epoch和step
|
||||
checkpoint = {
|
||||
'epoch': epoch,
|
||||
'step': step + 1,
|
||||
'optimizer_state_dict': optimizer.state_dict(),
|
||||
'scaler_state_dict': scaler.state_dict()
|
||||
}
|
||||
torch.save(checkpoint, f'{args.save_dir}/checkpoint.pth')
|
||||
model.train()
|
||||
|
||||
|
||||
|
|
@ -141,6 +152,7 @@ if __name__ == "__main__":
|
|||
parser.add_argument('--max_seq_len', default=512, type=int)
|
||||
parser.add_argument('--use_moe', default=False, type=bool)
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/sft_mini_512.jsonl")
|
||||
parser.add_argument("--resume", action="store_true", help="Resume training from checkpoint")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
|
@ -198,5 +210,20 @@ if __name__ == "__main__":
|
|||
model = DistributedDataParallel(model, device_ids=[ddp_local_rank])
|
||||
|
||||
iter_per_epoch = len(train_loader)
|
||||
for epoch in range(args.epochs):
|
||||
train_epoch(epoch, wandb)
|
||||
|
||||
start_epoch = 0
|
||||
start_step = 0
|
||||
if args.resume:
|
||||
checkpoint_path = f'{args.save_dir}/checkpoint.pth'
|
||||
if os.path.exists(checkpoint_path):
|
||||
checkpoint = torch.load(checkpoint_path)
|
||||
start_epoch = checkpoint['epoch']
|
||||
start_step = checkpoint['step']
|
||||
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
|
||||
scaler.load_state_dict(checkpoint['scaler_state_dict'])
|
||||
Logger(f"Resuming training from epoch {start_epoch + 1}, step {start_step}")
|
||||
else:
|
||||
Logger("Checkpoint not found, starting training from scratch.")
|
||||
|
||||
for epoch in range(start_epoch, args.epochs):
|
||||
train_epoch(epoch, wandb, start_step if epoch == start_epoch else 0)
|
||||
|
|
@ -28,11 +28,22 @@ def Logger(content):
|
|||
def get_lr(current_step, total_steps, lr):
|
||||
return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
|
||||
|
||||
|
||||
def train_epoch(epoch, wandb):
|
||||
loss_fct = nn.CrossEntropyLoss(reduction='none')
|
||||
|
||||
def train_epoch(epoch, start_step, wandb):
|
||||
start_time = time.time()
|
||||
if epoch == start_epoch:
|
||||
train_loader_iter = iter(train_loader)
|
||||
for _ in range(start_step + 1):
|
||||
next(train_loader_iter)
|
||||
for step, (X, Y, loss_mask) in enumerate(train_loader_iter, start=start_step + 1):
|
||||
train_step(epoch, step, X, Y, loss_mask, wandb, start_time)
|
||||
else:
|
||||
for step, (X, Y, loss_mask) in enumerate(train_loader):
|
||||
train_step(epoch, step, X, Y, loss_mask, wandb, start_time)
|
||||
|
||||
|
||||
def train_step(epoch, step, X, Y, loss_mask, wandb, start_time):
|
||||
X = X.to(args.device)
|
||||
Y = Y.to(args.device)
|
||||
loss_mask = loss_mask.to(args.device)
|
||||
|
|
@ -91,6 +102,17 @@ def train_epoch(epoch, wandb):
|
|||
|
||||
state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存
|
||||
torch.save(state_dict, ckp)
|
||||
|
||||
# 保存训练状态
|
||||
training_state = {
|
||||
'epoch': epoch,
|
||||
'step': step,
|
||||
'optimizer_state_dict': optimizer.state_dict(),
|
||||
'scaler_state_dict': scaler.state_dict() if scaler is not None else None
|
||||
}
|
||||
training_state_path = f'{args.save_dir}/training_state.pth'
|
||||
torch.save(training_state, training_state_path)
|
||||
|
||||
model.train()
|
||||
|
||||
|
||||
|
|
@ -115,28 +137,28 @@ def init_distributed_mode():
|
|||
|
||||
# torchrun --nproc_per_node 2 1-pretrain.py
|
||||
if __name__ == "__main__":
|
||||
# 命令行参数解析
|
||||
parser = argparse.ArgumentParser(description="MiniMind Pretraining")
|
||||
parser.add_argument("--out_dir", type=str, default="../out")
|
||||
# 若要以最快速度实现zero则epochs设置为1轮；否则应当利用有限的数据训练2~6个epochs。
|
||||
parser.add_argument("--epochs", type=int, default=1)
|
||||
parser.add_argument("--batch_size", type=int, default=32)
|
||||
parser.add_argument("--out_dir", type=str, default="../out") # 输出目录
|
||||
parser.add_argument("--epochs", type=int, default=8) # 训练轮次
|
||||
parser.add_argument("--batch_size", type=int, default=8) # 物理批次大小
|
||||
parser.add_argument("--learning_rate", type=float, default=5e-4)
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16")
|
||||
parser.add_argument("--use_wandb", action="store_true")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16") # 混合精度类型
|
||||
parser.add_argument("--use_wandb", action="store_true") # 是否使用WandB
|
||||
parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain")
|
||||
parser.add_argument("--num_workers", type=int, default=1)
|
||||
parser.add_argument("--ddp", action="store_true")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=8)
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0)
|
||||
parser.add_argument("--warmup_iters", type=int, default=0)
|
||||
parser.add_argument("--log_interval", type=int, default=100)
|
||||
parser.add_argument("--save_interval", type=int, default=100)
|
||||
parser.add_argument('--local_rank', type=int, default=-1)
|
||||
parser.add_argument('--hidden_size', default=512, type=int)
|
||||
parser.add_argument('--num_hidden_layers', default=8, type=int)
|
||||
parser.add_argument('--max_seq_len', default=512, type=int)
|
||||
parser.add_argument('--use_moe', default=False, type=bool)
|
||||
parser.add_argument("--num_workers", type=int, default=1) # 数据加载线程数
|
||||
parser.add_argument("--ddp", action="store_true") # 是否启用DDP
|
||||
parser.add_argument("--accumulation_steps", type=int, default=4) # 梯度累积步数
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0) # 梯度裁剪阈值
|
||||
parser.add_argument("--log_interval", type=int, default=100) # 日志间隔（步）
|
||||
parser.add_argument("--save_interval", type=int, default=100) # 保存间隔（步）
|
||||
parser.add_argument('--local_rank', type=int, default=-1) # DDP自动传入参数
|
||||
# 模型架构参数
|
||||
parser.add_argument('--hidden_size', default=512, type=int) # 隐藏层维度
|
||||
parser.add_argument('--num_hidden_layers', default=8, type=int) # 层数
|
||||
parser.add_argument('--max_seq_len', default=4096, type=int) # 序列长度
|
||||
parser.add_argument('--use_moe', default=False, type=bool) # 是否使用MoE
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/pretrain_hq.jsonl")
|
||||
args = parser.parse_args()
|
||||
|
||||
|
|
@ -194,5 +216,19 @@ if __name__ == "__main__":
|
|||
model = DistributedDataParallel(model, device_ids=[ddp_local_rank])
|
||||
|
||||
iter_per_epoch = len(train_loader)
|
||||
for epoch in range(args.epochs):
|
||||
train_epoch(epoch, wandb)
|
||||
|
||||
# 检查是否存在保存的训练状态文件
|
||||
training_state_path = f'{args.save_dir}/training_state.pth'
|
||||
if os.path.exists(training_state_path):
|
||||
training_state = torch.load(training_state_path)
|
||||
start_epoch = training_state['epoch']
|
||||
start_step = training_state['step']
|
||||
optimizer.load_state_dict(training_state['optimizer_state_dict'])
|
||||
if scaler is not None and training_state['scaler_state_dict'] is not None:
|
||||
scaler.load_state_dict(training_state['scaler_state_dict'])
|
||||
else:
|
||||
start_epoch = 0
|
||||
start_step = 0
|
||||
|
||||
for epoch in range(start_epoch, args.epochs):
|
||||
train_epoch(epoch, start_step if epoch == start_epoch else 0, wandb)
|
||||
Loading…
Reference in New Issue