441 lines
21 KiB
Python
441 lines
21 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
|
||
import torch.distributed as dist
|
||
from torch import optim, nn
|
||
from torch.nn.parallel import DistributedDataParallel
|
||
from torch.optim.lr_scheduler import CosineAnnealingLR
|
||
from torch.utils.data import DataLoader, DistributedSampler
|
||
# 移除通信分析工具导入
|
||
from contextlib import nullcontext
|
||
from typing import Optional
|
||
|
||
from transformers import AutoTokenizer
|
||
|
||
from model.model import MiniMindLM
|
||
from model.LMConfig import LMConfig
|
||
from model.dataset import PretrainDataset
|
||
|
||
warnings.filterwarnings('ignore')
|
||
|
||
|
||
def Logger(content):
|
||
# 如果没有使用ddp或者ddp的主设备，那么就打印
|
||
if not ddp or dist.get_rank() == 0:
|
||
print(content)
|
||
|
||
|
||
def get_lr(current_step, total_steps, lr):
|
||
# 更新学习率
|
||
# \text{get\_lr}(c, t, l) = \frac{l}{10} + 0.5 \cdot l \cdot \left(1 + \cos\left(\frac{\pi \cdot c}{t}\right)\right)
|
||
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')
|
||
start_time = time.time()
|
||
# 在函数开始处定义moe_path，避免在异常处理中引用未定义变量
|
||
moe_path = '_moe' if lm_config.use_moe else ''
|
||
|
||
# 添加CUDA事件来分析性能
|
||
if args.profile and (not ddp or dist.get_rank() == 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)
|
||
|
||
# 移除CUDA图优化代码
|
||
|
||
# 预取数据
|
||
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([t.to(args.device, non_blocking=True) for t in batch])
|
||
except StopIteration:
|
||
break
|
||
|
||
for step in range(len(train_loader)):
|
||
try:
|
||
# 计时数据加载
|
||
if args.profile and (not ddp or dist.get_rank() == 0):
|
||
data_start.record()
|
||
|
||
# 使用预取的数据
|
||
if prefetch_batches:
|
||
X, Y, loss_mask = prefetch_batches.pop(0)
|
||
else:
|
||
# 如果预取队列为空，直接加载
|
||
X, Y, loss_mask = [t.to(args.device) for t in next(data_iter)]
|
||
|
||
# 异步预取下一批数据
|
||
if step + prefetch_factor < len(train_loader):
|
||
try:
|
||
batch = next(data_iter)
|
||
prefetch_batches.append([t.to(args.device, non_blocking=True) for t in batch])
|
||
except StopIteration:
|
||
pass
|
||
|
||
if args.profile and (not ddp or dist.get_rank() == 0):
|
||
data_end.record()
|
||
|
||
# 更新学习率
|
||
lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)
|
||
for param_group in optimizer.param_groups:
|
||
param_group['lr'] = lr
|
||
|
||
# 计时前向传播
|
||
if args.profile and (not ddp or dist.get_rank() == 0):
|
||
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:
|
||
if hasattr(model, 'module'):
|
||
# DDP情况
|
||
aux_loss = sum(l.feed_forward.aux_loss for l in model.module.layers
|
||
if hasattr(l.feed_forward, 'aux_loss'))
|
||
else:
|
||
# 非DDP情况
|
||
aux_loss = sum(l.feed_forward.aux_loss for l in model.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
|
||
|
||
# 反向传播
|
||
scaler.scale(loss).backward()
|
||
|
||
if args.profile and (not ddp or dist.get_rank() == 0):
|
||
forward_end.record()
|
||
backward_start.record()
|
||
|
||
# Print data types for debugging
|
||
if step == 0 and (not ddp or dist.get_rank() == 0): # Print only for the first step of the first epoch on the main process
|
||
Logger("---- Data Type Check ----")
|
||
Logger(f"X.dtype: {X.dtype}")
|
||
if hasattr(model, 'module'): # DDP case
|
||
Logger(f"Model parameter dtype: {next(model.module.parameters()).dtype}")
|
||
else: # Non-DDP case
|
||
Logger(f"Model parameter dtype: {next(model.parameters()).dtype}")
|
||
Logger(f"res.logits.dtype: {res.logits.dtype}")
|
||
Logger(f"loss.dtype: {loss.dtype}")
|
||
Logger("-------------------------")
|
||
|
||
if args.profile and (not ddp or dist.get_rank() == 0):
|
||
backward_end.record()
|
||
|
||
# 在每一步都进行性能分析，而不仅仅是在梯度累积完成时
|
||
if (step + 1) % args.profile_interval == 0:
|
||
# 记录优化器时间（如果是梯度累积步骤）
|
||
if (step + 1) % args.accumulation_steps == 0:
|
||
optimizer_start.record()
|
||
|
||
# 优化器步骤
|
||
if (step + 1) % args.accumulation_steps == 0:
|
||
if args.profile and (not ddp or dist.get_rank() == 0):
|
||
if (step + 1) % args.profile_interval != 0:
|
||
optimizer_start.record()
|
||
|
||
scaler.unscale_(optimizer)
|
||
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||
|
||
scaler.step(optimizer)
|
||
scaler.update()
|
||
|
||
optimizer.zero_grad(set_to_none=True)
|
||
|
||
if args.profile and (not ddp or dist.get_rank() == 0):
|
||
optimizer_end.record()
|
||
|
||
# 性能分析输出（每profile_interval步）
|
||
if args.profile and (not ddp or dist.get_rank() == 0) and (step + 1) % args.profile_interval == 0:
|
||
# 同步CUDA事件以获取准确的计时
|
||
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)
|
||
|
||
# 只有在梯度累积步骤完成时才有优化器时间
|
||
if (step + 1) % args.accumulation_steps == 0:
|
||
optimizer_time = optimizer_start.elapsed_time(optimizer_end)
|
||
total_compute_time = forward_time + backward_time + optimizer_time
|
||
Logger(f"性能分析 - 步骤 {step+1}:")
|
||
Logger(f" 数据加载时间: {data_time:.2f} ms")
|
||
Logger(f" 前向传播时间: {forward_time:.2f} ms")
|
||
Logger(f" 反向传播时间: {backward_time:.2f} ms")
|
||
Logger(f" 优化器时间: {optimizer_time:.2f} ms")
|
||
Logger(f" 总计算时间: {total_compute_time:.2f} ms")
|
||
Logger(f" 计算/数据比例: {total_compute_time / data_time:.2f}")
|
||
else:
|
||
# 非梯度累积步骤，没有优化器时间
|
||
total_compute_time = forward_time + backward_time
|
||
Logger(f"性能分析 - 步骤 {step+1} (梯度累积中):")
|
||
Logger(f" 数据加载时间: {data_time:.2f} ms")
|
||
Logger(f" 前向传播时间: {forward_time:.2f} ms")
|
||
Logger(f" 反向传播时间: {backward_time:.2f} ms")
|
||
Logger(f" 总计算时间: {total_compute_time:.2f} ms")
|
||
Logger(f" 计算/数据比例: {total_compute_time / data_time:.2f}")
|
||
|
||
# 打印日志
|
||
if step % args.log_interval == 0:
|
||
spend_time = time.time() - start_time
|
||
Logger(
|
||
'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format(
|
||
epoch + 1,
|
||
args.epochs,
|
||
step,
|
||
iter_per_epoch,
|
||
loss.item() * args.accumulation_steps,
|
||
optimizer.param_groups[-1]['lr'],
|
||
spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
|
||
|
||
if (wandb is not None) and (not ddp or dist.get_rank() == 0):
|
||
log_dict = {
|
||
"loss": loss.item() * args.accumulation_steps,
|
||
"lr": optimizer.param_groups[-1]['lr'],
|
||
"epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60
|
||
}
|
||
|
||
# 如果启用了性能分析，也记录性能指标
|
||
if args.profile and (step + 1) % args.profile_interval == 0:
|
||
# 基本性能指标
|
||
perf_dict = {
|
||
"data_time_ms": data_time,
|
||
"forward_time_ms": forward_time,
|
||
"backward_time_ms": backward_time
|
||
}
|
||
|
||
# 只有在梯度累积步骤完成时才有优化器时间
|
||
if (step + 1) % args.accumulation_steps == 0:
|
||
total_compute_time = forward_time + backward_time + optimizer_time
|
||
perf_dict.update({
|
||
"optimizer_time_ms": optimizer_time,
|
||
"compute_time_ms": total_compute_time
|
||
})
|
||
else:
|
||
total_compute_time = forward_time + backward_time
|
||
perf_dict.update({
|
||
"compute_time_ms": total_compute_time
|
||
})
|
||
|
||
log_dict.update(perf_dict)
|
||
|
||
wandb.log(log_dict)
|
||
|
||
# 移除通信分析代码
|
||
|
||
# 保存模型
|
||
if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
|
||
model.eval()
|
||
# 使用函数开始处定义的moe_path变量
|
||
ckp = f'{args.save_dir}/pretrain_{lm_config.dim}{moe_path}.pth'
|
||
|
||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||
state_dict = model.module.state_dict() #获取模型参数
|
||
else:
|
||
state_dict = model.state_dict() #获取模型参数
|
||
|
||
torch.save(state_dict, ckp) #只保存参数
|
||
model.train()
|
||
|
||
except Exception as e:
|
||
print(f"Error occurred: {str(e)}")
|
||
save_path = f'{args.save_dir}/pretrain_{lm_config.dim}{moe_path}_nanERROR.pth'
|
||
if os.path.exists(save_path):
|
||
os.remove(save_path)
|
||
|
||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||
state_dict = model.module.state_dict()
|
||
else:
|
||
state_dict = model.state_dict()
|
||
torch.save(state_dict, save_path)
|
||
|
||
for name, param in model.named_parameters():
|
||
if param.grad is not None and torch.isnan(param.grad).any():
|
||
print(f"NaN gradient in parameter: {name}")
|
||
|
||
for name, param in model.named_parameters():
|
||
if param.grad is not None and torch.isnan(param.grad).any():
|
||
print(f"Parameter {name} values: {param.data}")
|
||
print(f"Parameter {name} gradients: {param.grad}")
|
||
|
||
raise ValueError("NaN gradient detected")
|
||
|
||
|
||
def init_model(lm_config, pretrained_embedding_path: Optional[str] = None):
|
||
# 加载tokenizer
|
||
tokenizer = AutoTokenizer.from_pretrained('/mnt/lzn/Minimind/Minimind/model/minimind_tokenizer')
|
||
# 加载模型
|
||
model = MiniMindLM(lm_config).to(args.device)
|
||
|
||
# Load pretrained token embeddings if path is provided
|
||
if pretrained_embedding_path and os.path.exists(pretrained_embedding_path):
|
||
Logger(f"Loading pretrained token embeddings from {pretrained_embedding_path}")
|
||
embedding_weights = torch.load(pretrained_embedding_path, map_location=args.device)
|
||
model.tok_embeddings.load_state_dict(embedding_weights)
|
||
Logger("Successfully loaded pretrained token embeddings.")
|
||
elif pretrained_embedding_path:
|
||
Logger(f"Warning: Pretrained embedding path {pretrained_embedding_path} provided but file does not exist. Initializing embeddings from scratch.")
|
||
|
||
# 打印模型参数
|
||
Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||
return model, tokenizer
|
||
|
||
|
||
# 移除通信分析函数
|
||
|
||
|
||
def init_distributed_mode():
|
||
if not ddp: return #如果没有启用分布式数据并行(DDP)，直接返回，不执行任何操作。
|
||
global ddp_local_rank, DEVICE #声明这两个变量为全局变量，以便在函数外部也能访问它们。
|
||
|
||
dist.init_process_group(backend="nccl") #初始化分布式进程组，使用NCCL后端（NVIDIA Collective Communications Library），这是NVIDIA GPU之间通信的优化库。
|
||
ddp_rank = int(os.environ["RANK"]) #从环境变量获取当前进程的全局编号。
|
||
ddp_local_rank = int(os.environ["LOCAL_RANK"]) #从环境变量获取当前进程的本地编号。
|
||
ddp_world_size = int(os.environ["WORLD_SIZE"]) #从环境变量获取当前进程组中的进程总数。
|
||
DEVICE = f"cuda:{ddp_local_rank}" #根据本地编号选择GPU设备。
|
||
torch.cuda.set_device(DEVICE) #设置当前进程的GPU设备。
|
||
|
||
|
||
# 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=3)
|
||
parser.add_argument("--batch_size", type=int, default=24)
|
||
parser.add_argument("--learning_rate", type=float, default=2e-4)
|
||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu") #如果GPU可用，则使用GPU，否则使用CPU。
|
||
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("--ddp", action="store_true")
|
||
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('--local_rank', type=int, default=-1) #本地进程编号，用于分布式训练。
|
||
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) #是否使用MOE，用于控制是否使用MOE。
|
||
parser.add_argument('--disable_db', action='store_true', help="禁用数据库功能，使用固定值1e-4替代") #禁用数据库功能，启用特殊模式
|
||
parser.add_argument("--data_path", type=str, default="/mnt/lzn/Minimind/dataset/dir/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")
|
||
args = parser.parse_args()
|
||
print(args)
|
||
|
||
|
||
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 # 添加FlashAttention支持
|
||
) #创建LMConfig对象，用于控制模型配置。
|
||
args.save_dir = os.path.join(args.out_dir) #创建保存目录。
|
||
os.makedirs(args.save_dir, exist_ok=True) #创建保存目录。
|
||
os.makedirs(args.out_dir, exist_ok=True) #创建输出目录。
|
||
tokens_per_iter = args.batch_size * lm_config.max_seq_len #计算每个迭代步骤的token数量。
|
||
print(f"tokens_per_iter: {tokens_per_iter}")
|
||
device_type = "cuda" if "cuda" in args.device else "cpu" #确定设备类型。
|
||
|
||
# Determine the torch dtype
|
||
pt_dtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[args.dtype]
|
||
|
||
args.wandb_run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
|
||
|
||
ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=pt_dtype)
|
||
|
||
ddp = int(os.environ.get("RANK", -1)) != -1 # is this a ddp run?
|
||
ddp_local_rank, DEVICE = 0, "cuda:0"
|
||
|
||
base_seed = 1337
|
||
torch.manual_seed(base_seed)
|
||
torch.cuda.manual_seed(base_seed)
|
||
|
||
if ddp:
|
||
init_distributed_mode()
|
||
args.device = torch.device(DEVICE)
|
||
rank = dist.get_rank()
|
||
torch.manual_seed(base_seed + rank)
|
||
# 同时设置 CUDA 的随机种子
|
||
torch.cuda.manual_seed(base_seed + rank)
|
||
|
||
if args.use_wandb and (not ddp or ddp_local_rank == 0):
|
||
import wandb
|
||
|
||
# Merge args and lm_config parameters for wandb config
|
||
config = vars(args).copy()
|
||
config.update(lm_config.__dict__)
|
||
|
||
wandb.init(project=args.wandb_project, name=args.wandb_run_name, config=config)
|
||
else:
|
||
wandb = None
|
||
model, tokenizer = init_model(lm_config, args.pretrained_embedding_path)
|
||
train_ds = PretrainDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
|
||
train_sampler = DistributedSampler(train_ds) if ddp else None
|
||
# 优化DataLoader配置
|
||
train_loader = DataLoader(
|
||
train_ds,
|
||
batch_size=args.batch_size,
|
||
pin_memory=True,
|
||
pin_memory_device=f"cuda:{ddp_local_rank}" if ddp else "cuda:0", # 指定pin_memory设备
|
||
drop_last=False,
|
||
shuffle=False,
|
||
num_workers=args.num_workers,
|
||
sampler=train_sampler,
|
||
persistent_workers=True if args.num_workers > 0 else False, # 保持worker进程活跃
|
||
prefetch_factor=2 if args.num_workers > 0 else None # 预取因子
|
||
)
|
||
|
||
# 只有在使用float16时才启用GradScaler，bfloat16不需要
|
||
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
|
||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||
|
||
if ddp:
|
||
model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
|
||
# 保留find_unused_parameters=True参数，因为模型中确实有未使用的参数
|
||
model = DistributedDataParallel(model, device_ids=[ddp_local_rank], find_unused_parameters=True)
|
||
|
||
# 暂时保留set_detect_anomaly以便调试
|
||
# 训练稳定后可以注释掉这行来提高速度
|
||
torch.autograd.set_detect_anomaly(True)
|
||
iter_per_epoch = len(train_loader)
|
||
for epoch in range(args.epochs):
|
||
train_epoch(epoch, wandb)
|