399 lines
18 KiB
Python
399 lines
18 KiB
Python
import os
|
||
# 设置环境变量
|
||
os.environ["WANDB_MODE"] = "offline" # 或者使用 "dryrun"
|
||
import platform
|
||
import argparse
|
||
import time
|
||
import math
|
||
import warnings
|
||
import pandas as pd
|
||
import torch
|
||
from torch import optim, nn
|
||
from torch.utils.data import DataLoader
|
||
from contextlib import nullcontext
|
||
from typing import Optional
|
||
import datetime # Add datetime for time formatting
|
||
from accelerate import Accelerator
|
||
from accelerate.utils import set_seed
|
||
from accelerate.utils import DeepSpeedPlugin
|
||
from accelerate.utils import DistributedDataParallelKwargs
|
||
from transformers import AutoTokenizer, get_cosine_schedule_with_warmup
|
||
|
||
from model.model import MiniMindLM
|
||
from model.LMConfig import LMConfig
|
||
from model.dataset import PretrainDataset
|
||
|
||
warnings.filterwarnings('ignore')
|
||
|
||
# 日志记录函数
|
||
def Logger(msg, accelerator=None):
|
||
# 如果没有提供accelerator，则只在主进程打印
|
||
if accelerator is None or accelerator.is_main_process:
|
||
print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {msg}")
|
||
|
||
# Helper function to format seconds into HH:MM:SS
|
||
def format_time(seconds):
|
||
return str(datetime.timedelta(seconds=int(seconds)))
|
||
|
||
# 获取学习率函数
|
||
def get_lr(it, num_iters, learning_rate):
|
||
# 余弦学习率衰减
|
||
return learning_rate * 0.5 * (1.0 + math.cos(math.pi * it / num_iters))
|
||
|
||
# 初始化模型函数
|
||
def init_model(lm_config, pretrained_embedding_path=None):
|
||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||
model = MiniMindLM(lm_config)
|
||
|
||
# 如果提供了预训练的嵌入权重，加载它们
|
||
if pretrained_embedding_path:
|
||
Logger(f"Loading pretrained token embeddings from {pretrained_embedding_path}")
|
||
pretrained_embeddings = torch.load(pretrained_embedding_path)
|
||
model.tok_embeddings.weight.data.copy_(pretrained_embeddings)
|
||
model.output.weight.data.copy_(pretrained_embeddings) # 共享权重
|
||
|
||
Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||
return model, tokenizer
|
||
|
||
def train_epoch(epoch, accelerator, model, train_loader, optimizer, scheduler, args, ctx, overall_start_time):
|
||
loss_fct = nn.CrossEntropyLoss(reduction='none')
|
||
epoch_start_time = time.time()
|
||
total_steps_in_epoch = len(train_loader)
|
||
total_training_steps = args.epochs * total_steps_in_epoch
|
||
moe_path = '_moe' if args.use_moe else ''
|
||
|
||
# 添加CUDA事件来分析性能 (只在主进程进行)
|
||
if args.profile and accelerator.is_main_process:
|
||
data_start = torch.cuda.Event(enable_timing=True)
|
||
data_end = torch.cuda.Event(enable_timing=True)
|
||
forward_start = torch.cuda.Event(enable_timing=True)
|
||
forward_end = torch.cuda.Event(enable_timing=True)
|
||
backward_start = torch.cuda.Event(enable_timing=True)
|
||
backward_end = torch.cuda.Event(enable_timing=True)
|
||
optimizer_start = torch.cuda.Event(enable_timing=True)
|
||
optimizer_end = torch.cuda.Event(enable_timing=True)
|
||
|
||
# 预取数据
|
||
prefetch_factor = 2 # 预取的批次数
|
||
data_iter = iter(train_loader)
|
||
prefetch_batches = []
|
||
|
||
# 预取初始批次
|
||
for _ in range(min(prefetch_factor, len(train_loader))):
|
||
try:
|
||
batch = next(data_iter)
|
||
prefetch_batches.append(batch)
|
||
except StopIteration:
|
||
break
|
||
|
||
# 在开始循环前初始化日志记录所需变量
|
||
last_log_time = epoch_start_time
|
||
|
||
for step in range(total_steps_in_epoch):
|
||
try:
|
||
# 计时数据加载 (只在主进程进行)
|
||
if args.profile and accelerator.is_main_process:
|
||
data_start.record()
|
||
|
||
# 使用预取的数据
|
||
if prefetch_batches:
|
||
X, Y, loss_mask = prefetch_batches.pop(0)
|
||
else:
|
||
# 如果预取队列为空，直接加载
|
||
X, Y, loss_mask = next(data_iter)
|
||
|
||
# 异步预取下一批数据
|
||
if step + prefetch_factor < len(train_loader):
|
||
try:
|
||
batch = next(data_iter)
|
||
prefetch_batches.append(batch)
|
||
except StopIteration:
|
||
pass
|
||
|
||
# 计时数据加载结束 (只在主进程进行)
|
||
if args.profile and accelerator.is_main_process:
|
||
data_end.record()
|
||
|
||
# 更新学习率
|
||
if scheduler is not None:
|
||
scheduler.step()
|
||
|
||
# 计时前向传播 (只在主进程进行)
|
||
if args.profile and accelerator.is_main_process:
|
||
forward_start.record()
|
||
|
||
# 前向传播
|
||
with ctx:
|
||
res = model(X)
|
||
loss = loss_fct(
|
||
res.logits.view(-1, res.logits.size(-1)),
|
||
Y.view(-1)
|
||
).view(Y.size())
|
||
loss = (loss * loss_mask).sum() / loss_mask.sum()
|
||
# 添加辅助损失，如果存在的话
|
||
try:
|
||
aux_loss = sum(l.feed_forward.aux_loss for l in model.module.layers
|
||
if hasattr(l.feed_forward, 'aux_loss'))
|
||
loss += aux_loss
|
||
except Exception as e:
|
||
Logger(f"Warning: Could not add auxiliary loss: {e}")
|
||
# 如果出错，不添加辅助损失
|
||
loss = loss / args.accumulation_steps
|
||
|
||
# 计时前向传播结束 (只在主进程进行)
|
||
if args.profile and accelerator.is_main_process:
|
||
forward_end.record()
|
||
|
||
# 计时反向传播 (只在主进程进行)
|
||
if args.profile and accelerator.is_main_process:
|
||
backward_start.record()
|
||
|
||
# 反向传播
|
||
# 当使用DeepSpeed时，它会自动处理梯度累积和梯度裁剪
|
||
accelerator.backward(loss)
|
||
|
||
# 计时反向传播结束 (只在主进程进行)
|
||
if args.profile and accelerator.is_main_process:
|
||
backward_end.record()
|
||
|
||
# 计时优化器步骤 (只在主进程进行)
|
||
if args.profile and accelerator.is_main_process:
|
||
optimizer_start.record()
|
||
|
||
# 优化器步骤 - 当使用DeepSpeed时，它会自动处理梯度累积和梯度裁剪
|
||
# 只有在达到累积步数时才会执行优化器步骤
|
||
# 注意：当使用DeepSpeed时，它会自动处理梯度累积，所以我们不需要检查step % accumulation_steps
|
||
optimizer.step()
|
||
|
||
# 当使用DeepSpeed时，zero_grad()会在step()之后自动调用
|
||
# 但为了安全起见，我们仍然显式调用它
|
||
optimizer.zero_grad()
|
||
|
||
# 计时优化器步骤结束 (只在主进程进行)
|
||
if args.profile and accelerator.is_main_process:
|
||
optimizer_end.record()
|
||
|
||
# 打印训练信息 (只在主进程进行)
|
||
if (step + 1) % args.log_interval == 0 and accelerator.is_main_process:
|
||
current_time = time.time()
|
||
# 计算性能指标
|
||
if args.profile:
|
||
torch.cuda.synchronize()
|
||
# 使用自上次日志以来的时间计算性能指标，而不是总时间
|
||
data_time = data_start.elapsed_time(data_end)
|
||
forward_time = forward_start.elapsed_time(forward_end)
|
||
backward_time = backward_start.elapsed_time(backward_end)
|
||
optimizer_time = optimizer_start.elapsed_time(optimizer_end)
|
||
iter_time = (current_time - last_log_time) * 1000 / args.log_interval # avg ms per iteration since last log
|
||
# total_time_ms = data_time + forward_time + backward_time + optimizer_time
|
||
|
||
# 打印性能分析
|
||
if (step + 1) % (args.log_interval * args.profile_interval) == 0:
|
||
Logger(f"性能分析 (Avg/iter over last {args.log_interval} steps) - "
|
||
f"Data: {data_time/args.log_interval:.2f}ms, "
|
||
f"Fwd: {forward_time/args.log_interval:.2f}ms, "
|
||
f"Bwd: {backward_time/args.log_interval:.2f}ms, "
|
||
f"Optim: {optimizer_time/args.log_interval:.2f}ms, "
|
||
f"Iter Time: {iter_time:.2f}ms", accelerator)
|
||
# 重置事件以便下次测量从0开始
|
||
data_start = torch.cuda.Event(enable_timing=True)
|
||
data_end = torch.cuda.Event(enable_timing=True)
|
||
forward_start = torch.cuda.Event(enable_timing=True)
|
||
forward_end = torch.cuda.Event(enable_timing=True)
|
||
backward_start = torch.cuda.Event(enable_timing=True)
|
||
backward_end = torch.cuda.Event(enable_timing=True)
|
||
optimizer_start = torch.cuda.Event(enable_timing=True)
|
||
optimizer_end = torch.cuda.Event(enable_timing=True)
|
||
|
||
|
||
# 计算当前学习率
|
||
current_lr = optimizer.param_groups[0]['lr']
|
||
|
||
# 计算时间
|
||
epoch_elapsed_time = current_time - epoch_start_time
|
||
epoch_steps_done = step + 1
|
||
epoch_avg_step_time = epoch_elapsed_time / epoch_steps_done
|
||
epoch_remaining_time = epoch_avg_step_time * (total_steps_in_epoch - epoch_steps_done)
|
||
|
||
total_elapsed_time = current_time - overall_start_time
|
||
total_steps_done = epoch * total_steps_in_epoch + epoch_steps_done
|
||
total_avg_step_time = total_elapsed_time / total_steps_done if total_steps_done > 0 else 0
|
||
total_remaining_time = total_avg_step_time * (total_training_steps - total_steps_done) if total_steps_done > 0 else 0
|
||
|
||
# 计算训练速度 (基于最近的log_interval)
|
||
interval_elapsed_time = current_time - last_log_time
|
||
tokens_processed_interval = args.log_interval * args.batch_size * args.max_seq_len
|
||
tokens_per_sec = tokens_processed_interval / interval_elapsed_time if interval_elapsed_time > 0 else 0
|
||
last_log_time = current_time # 更新上次日志时间
|
||
|
||
Logger(f"Epoch {epoch+1}/{args.epochs}, Step {step+1}/{total_steps_in_epoch}, "
|
||
f"Loss: {loss.item()*args.accumulation_steps:.4f}, "
|
||
f"LR: {current_lr:.6f}, "
|
||
f"Speed: {tokens_per_sec:.2f} tokens/sec | "
|
||
f"Epoch Time Left: {format_time(epoch_remaining_time)} | "
|
||
f"Total Time Left: {format_time(total_remaining_time)}", accelerator)
|
||
|
||
# 保存模型 (只在主进程进行)
|
||
if (step + 1) % args.save_interval == 0 and accelerator.is_main_process:
|
||
# 使用函数开始处定义的moe_path变量
|
||
ckp = f'{args.save_dir}/pretrain_{args.dim}{moe_path}.pth'
|
||
|
||
# 获取解包后的模型
|
||
unwrapped_model = accelerator.unwrap_model(model)
|
||
|
||
# 保存模型参数
|
||
accelerator.save(unwrapped_model.state_dict(), ckp)
|
||
Logger(f"Model saved to {ckp}", accelerator)
|
||
|
||
except Exception as e:
|
||
Logger(f"Error in training step: {e}", accelerator)
|
||
import traceback
|
||
Logger(traceback.format_exc(), accelerator)
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="MiniMind Pretraining with Accelerate")
|
||
parser.add_argument("--out_dir", type=str, default="out")
|
||
parser.add_argument("--epochs", type=int, default=3)
|
||
parser.add_argument("--batch_size", type=int, default=24)
|
||
parser.add_argument("--learning_rate", type=float, default=2e-4)
|
||
parser.add_argument("--dtype", type=str, default="bfloat16")
|
||
parser.add_argument("--use_wandb", default=True, action="store_true")
|
||
parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain")
|
||
parser.add_argument("--num_workers", type=int, default=48)
|
||
parser.add_argument("--accumulation_steps", type=int, default=32)
|
||
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=10000)
|
||
parser.add_argument('--dim', default=1024, type=int)
|
||
parser.add_argument('--n_layers', default=32, type=int)
|
||
parser.add_argument('--max_seq_len', default=1024, type=int)
|
||
parser.add_argument('--use_moe', default=False, type=bool)
|
||
parser.add_argument('--disable_db', action='store_true', help="禁用数据库功能，使用固定值1e-4替代")
|
||
parser.add_argument("--data_path", type=str, default="./dataset/pretrain_hq.jsonl")
|
||
parser.add_argument("--pretrained_embedding_path", type=str, default=None, help="Path to pretrained token embedding weights (.pth file)")
|
||
parser.add_argument("--profile", action="store_true", default=True, help="启用性能分析")
|
||
parser.add_argument("--profile_interval", type=int, default=10, help="性能分析打印间隔（步数）")
|
||
parser.add_argument("--use_flash_attn", action="store_true", default=True, help="启用FlashAttention")
|
||
parser.add_argument("--knowlwdge_num", type=int, default=64*64,help="知识库的数据数目")
|
||
parser.add_argument("--knowlwdge_length", type=int, default=8,help="知识库的句子长度")
|
||
args = parser.parse_args()
|
||
|
||
# 初始化accelerator
|
||
# 设置ddp_kwargs以处理未使用的参数
|
||
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
|
||
# 创建DeepSpeedPlugin对象
|
||
ds_plugin = DeepSpeedPlugin(
|
||
gradient_accumulation_steps=args.accumulation_steps,
|
||
gradient_clipping=args.grad_clip,
|
||
zero_stage=2, # 使用ZeRO-2优化
|
||
offload_optimizer_device="cpu", # 将优化器状态卸载到CPU
|
||
offload_param_device="none", # 不将参数卸载到CPU
|
||
)
|
||
accelerator = Accelerator(
|
||
kwargs_handlers=[ddp_kwargs],
|
||
deepspeed_plugin=ds_plugin,
|
||
mixed_precision="bf16" if args.dtype == "bfloat16" else "fp16" if args.dtype == "float16" else "no"
|
||
)
|
||
|
||
# 设置随机种子
|
||
set_seed(1337 + accelerator.process_index)
|
||
|
||
# 配置模型
|
||
lm_config = LMConfig(
|
||
dim=args.dim,
|
||
n_layers=args.n_layers,
|
||
max_seq_len=args.max_seq_len,
|
||
use_moe=args.use_moe,
|
||
disable_db=args.disable_db,
|
||
flash_attn=args.use_flash_attn,
|
||
knowlwdge_num=args.knowlwdge_num,
|
||
knowlwdge_length=args.knowlwdge_length
|
||
)
|
||
|
||
# 创建保存目录
|
||
args.save_dir = os.path.join(args.out_dir)
|
||
if accelerator.is_main_process:
|
||
os.makedirs(args.save_dir, exist_ok=True)
|
||
os.makedirs(args.out_dir, exist_ok=True)
|
||
|
||
# 计算每次迭代的token数量
|
||
tokens_per_iter = args.batch_size * lm_config.max_seq_len
|
||
Logger(f"tokens_per_iter: {tokens_per_iter}", accelerator)
|
||
|
||
# 设置数据类型
|
||
pt_dtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[args.dtype]
|
||
|
||
# 设置wandb运行名称
|
||
args.wandb_run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
|
||
|
||
# 设置自动混合精度上下文
|
||
ctx = nullcontext() if accelerator.device.type == "cpu" else torch.cuda.amp.autocast(dtype=pt_dtype)
|
||
|
||
# 初始化模型和tokenizer
|
||
model, tokenizer = init_model(lm_config, args.pretrained_embedding_path)
|
||
# 将accelerator传递给init_model函数中的Logger调用
|
||
Logger(f'模型初始化完成', accelerator)
|
||
|
||
# 处理位置编码张量问题
|
||
# 我们已经将复数版本的pos_cis替换为实数版本的pos_cis_real
|
||
# 但为了安全起见，我们仍然将其设置为不参与分布式训练
|
||
if hasattr(model, "pos_cis_real"):
|
||
Logger(f'检测到pos_cis_real实数张量，将其设置为不参与分布式训练', accelerator)
|
||
# 设置模型的_ddp_params_and_buffers_to_ignore属性
|
||
model._ddp_params_and_buffers_to_ignore = {"pos_cis_real"}
|
||
# 兼容旧版本，检查是否仍有pos_cis
|
||
elif hasattr(model, "pos_cis"):
|
||
Logger(f'检测到pos_cis复数张量，将其设置为不参与分布式训练', accelerator)
|
||
# 设置模型的_ddp_params_and_buffers_to_ignore属性
|
||
model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
|
||
|
||
# 创建数据集和数据加载器
|
||
train_ds = PretrainDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
|
||
train_loader = DataLoader(
|
||
train_ds,
|
||
batch_size=args.batch_size,
|
||
pin_memory=True,
|
||
drop_last=False,
|
||
shuffle=True,
|
||
num_workers=args.num_workers,
|
||
persistent_workers=True if args.num_workers > 0 else False,
|
||
prefetch_factor=2 if args.num_workers > 0 else None
|
||
)
|
||
|
||
# 创建优化器
|
||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||
|
||
# 创建学习率调度器
|
||
total_steps = len(train_loader) * args.epochs
|
||
warmup_steps = args.warmup_iters if args.warmup_iters > 0 else int(0.1 * total_steps)
|
||
scheduler = get_cosine_schedule_with_warmup(
|
||
optimizer,
|
||
num_warmup_steps=warmup_steps,
|
||
num_training_steps=total_steps
|
||
)
|
||
|
||
# 准备训练
|
||
model, optimizer, train_loader, scheduler = accelerator.prepare(
|
||
model, optimizer, train_loader, scheduler
|
||
)
|
||
|
||
# 初始化wandb
|
||
if args.use_wandb and accelerator.is_main_process:
|
||
import wandb
|
||
wandb.init(project=args.wandb_project, name=args.wandb_run_name, config=args)
|
||
else:
|
||
wandb = None
|
||
|
||
# 训练循环
|
||
overall_start_time = time.time() # Record overall start time
|
||
for epoch in range(args.epochs):
|
||
train_epoch(epoch, accelerator, model, train_loader, optimizer, scheduler, args, ctx, overall_start_time) # Pass overall start time
|
||
|
||
# 关闭wandb
|
||
if args.use_wandb and accelerator.is_main_process:
|
||
wandb.finish()
|
||
|
||
if __name__ == "__main__":
|
||
main()
|