791 lines
38 KiB
Python
791 lines
38 KiB
Python
import os
|
||
# 设置环境变量 - 将wandb替换为SwanLab
|
||
# os.environ["SWANLAB_MODE"] = "online" # SwanLab使用在线模式
|
||
import platform
|
||
import argparse
|
||
from tqdm import tqdm
|
||
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
|
||
import numpy as np
|
||
from sklearn.metrics.pairwise import cosine_similarity
|
||
import swanlab # 替换wandb导入
|
||
import gc # 添加垃圾回收模块
|
||
import psutil # 添加系统资源监控模块
|
||
|
||
from model.model import MiniMindLM, RMSNorm
|
||
from model.LMConfig import LMConfig
|
||
from model.dataset import PretrainDataset
|
||
|
||
warnings.filterwarnings('ignore')
|
||
|
||
# 内存监控辅助函数
|
||
def get_memory_usage():
|
||
"""获取当前内存使用情况"""
|
||
process = psutil.Process()
|
||
memory_info = process.memory_info()
|
||
return {
|
||
'rss_mb': memory_info.rss / 1024 / 1024, # 物理内存使用量（MB）
|
||
'vms_mb': memory_info.vms / 1024 / 1024, # 虚拟内存使用量（MB）
|
||
}
|
||
|
||
def get_cuda_memory_usage():
|
||
"""获取CUDA内存使用情况"""
|
||
if torch.cuda.is_available():
|
||
return {
|
||
'cuda_allocated_mb': torch.cuda.memory_allocated() / 1024 / 1024,
|
||
'cuda_reserved_mb': torch.cuda.memory_reserved() / 1024 / 1024,
|
||
'cuda_max_allocated_mb': torch.cuda.max_memory_allocated() / 1024 / 1024,
|
||
}
|
||
return {}
|
||
|
||
def get_tensor_memory_size(tensor_list):
|
||
"""计算tensor列表的总内存占用（MB）"""
|
||
total_size = 0
|
||
for batch in tensor_list:
|
||
if isinstance(batch, (list, tuple)):
|
||
for tensor in batch:
|
||
if isinstance(tensor, torch.Tensor):
|
||
total_size += tensor.numel() * tensor.element_size()
|
||
elif isinstance(batch, torch.Tensor):
|
||
total_size += batch.numel() * batch.element_size()
|
||
return total_size / 1024 / 1024 # 转换为MB
|
||
|
||
def log_memory_status(step, prefetch_batches, accelerator, stage="", detailed=False):
|
||
"""记录内存状态"""
|
||
if not accelerator.is_main_process:
|
||
return
|
||
|
||
memory_info = get_memory_usage()
|
||
cuda_info = get_cuda_memory_usage()
|
||
prefetch_memory = get_tensor_memory_size(prefetch_batches)
|
||
|
||
log_msg = f"[Memory Monitor] Step {step} {stage} - "
|
||
log_msg += f"Prefetch batches: {len(prefetch_batches)}, "
|
||
log_msg += f"Prefetch memory: {prefetch_memory:.2f}MB, "
|
||
log_msg += f"System RSS: {memory_info['rss_mb']:.2f}MB"
|
||
|
||
if cuda_info:
|
||
log_msg += f", CUDA allocated: {cuda_info['cuda_allocated_mb']:.2f}MB"
|
||
log_msg += f", CUDA reserved: {cuda_info['cuda_reserved_mb']:.2f}MB"
|
||
|
||
if detailed:
|
||
log_msg += f", System VMS: {memory_info['vms_mb']:.2f}MB"
|
||
if cuda_info:
|
||
log_msg += f", CUDA max allocated: {cuda_info['cuda_max_allocated_mb']:.2f}MB"
|
||
|
||
Logger(log_msg, accelerator)
|
||
|
||
# 日志记录函数
|
||
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, database_init_path=None, args=None):
|
||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||
model = MiniMindLM(lm_config)
|
||
|
||
# 默认模型初始化
|
||
Logger("Performing default model initialization...")
|
||
|
||
# 初始化嵌入层权重
|
||
nn.init.normal_(model.tok_embeddings.weight, mean=0.0, std=0.02)
|
||
|
||
# 初始化输出层权重（如果不共享权重的话）
|
||
if not hasattr(model.tok_embeddings, 'weight') or model.output.weight is not model.tok_embeddings.weight:
|
||
nn.init.normal_(model.output.weight, mean=0.0, std=0.02)
|
||
|
||
# 初始化所有线性层
|
||
for name, module in model.named_modules():
|
||
if isinstance(module, nn.Linear):
|
||
# 使用Xavier/Glorot初始化
|
||
nn.init.xavier_uniform_(module.weight)
|
||
if module.bias is not None:
|
||
nn.init.zeros_(module.bias)
|
||
elif isinstance(module, nn.Embedding):
|
||
# 嵌入层使用正态分布初始化
|
||
nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
||
elif isinstance(module, RMSNorm):
|
||
# RMSNorm的权重初始化为1
|
||
if hasattr(module, 'weight'):
|
||
nn.init.ones_(module.weight)
|
||
|
||
# 初始化位置编码相关参数
|
||
if hasattr(model.knowledge_dataset, 'keys'):
|
||
nn.init.normal_(model.knowledge_dataset.keys, mean=0.0, std=0.02)
|
||
|
||
Logger("Default model initialization completed")
|
||
|
||
# 如果提供了预训练的嵌入权重，加载它们
|
||
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) # 共享权重
|
||
|
||
if database_init_path:
|
||
import json
|
||
import os
|
||
|
||
# 数据库参数
|
||
knowledge_num = args.knowledge_num
|
||
knowledge_length = args.knowledge_length
|
||
|
||
# 检查是否使用缓存
|
||
cache_dir = os.path.dirname(args.cluster_cache_path)
|
||
if cache_dir:
|
||
os.makedirs(cache_dir, exist_ok=True)
|
||
|
||
processed_tensor = None
|
||
|
||
# 尝试加载缓存的处理结果
|
||
if not args.recompute_clusters and os.path.exists(args.cluster_cache_path):
|
||
try:
|
||
Logger(f"Loading cached processed results from {args.cluster_cache_path}")
|
||
processed_tensor = torch.load(args.cluster_cache_path)
|
||
|
||
# 验证缓存文件的形状是否可用
|
||
cached_knowledge_num, cached_knowledge_length = processed_tensor.shape
|
||
|
||
if cached_knowledge_length == knowledge_length:
|
||
if cached_knowledge_num >= knowledge_num:
|
||
# 缓存足够大，可以截取使用
|
||
processed_tensor = processed_tensor[:knowledge_num, :]
|
||
Logger(f"Successfully loaded cached data with shape {processed_tensor.shape}")
|
||
Logger(f"Truncated from cached shape ({cached_knowledge_num}, {cached_knowledge_length}) to required shape ({knowledge_num}, {knowledge_length})")
|
||
Logger("Skipping database initialization - using cached results")
|
||
else:
|
||
# 缓存太小，需要重新计算
|
||
Logger(f"Cached knowledge_num ({cached_knowledge_num}) < required knowledge_num ({knowledge_num}), recomputing...")
|
||
processed_tensor = None
|
||
else:
|
||
# knowledge_length不匹配，需要重新计算
|
||
Logger(f"Cached knowledge_length ({cached_knowledge_length}) != required knowledge_length ({knowledge_length}), recomputing...")
|
||
processed_tensor = None
|
||
except Exception as e:
|
||
Logger(f"Failed to load cached data: {e}, recomputing...")
|
||
processed_tensor = None
|
||
|
||
# 只有在没有有效缓存时才进行数据库初始化和处理
|
||
if processed_tensor is None:
|
||
Logger(f"Loading database initialization data from {database_init_path}")
|
||
|
||
# 1. 加载JSON文件
|
||
with open(database_init_path, 'r', encoding='utf-8') as f:
|
||
database_data = json.load(f)
|
||
|
||
# 提取sentences列表
|
||
sentences_data = database_data.get('sentences', [])
|
||
Logger(f"Loaded {len(sentences_data)} sentences from database")
|
||
|
||
# 2. 按照importance_score进行排序（从高到低）
|
||
sorted_sentences = sorted(sentences_data, key=lambda x: x.get('importance_score', 0.0), reverse=True)
|
||
Logger(f"Sorted sentences by importance score (highest: {sorted_sentences[0].get('importance_score', 0.0)}, lowest: {sorted_sentences[-1].get('importance_score', 0.0)})")
|
||
|
||
# 3. 处理每条数据，不进行聚类
|
||
Logger("Processing individual sentences...")
|
||
processed_rows = []
|
||
|
||
# 获取空token的id（用于填充）
|
||
pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
|
||
|
||
# 处理所需数量的句子
|
||
num_to_process = min(knowledge_num, len(sorted_sentences))
|
||
|
||
for i in range(num_to_process):
|
||
sentence_data = sorted_sentences[i]
|
||
sentence = sentence_data.get('corrected_sentence', '')
|
||
|
||
# 将句子转换为tokens
|
||
sentence_tokens = tokenizer.encode(sentence, add_special_tokens=False)
|
||
|
||
# 截断或填充到knowledge_length
|
||
if len(sentence_tokens) > knowledge_length:
|
||
# 如果超过长度，截断
|
||
sentence_tokens = sentence_tokens[:knowledge_length]
|
||
Logger(f"Sentence {i+1} truncated from {len(tokenizer.encode(sentence, add_special_tokens=False))} to {knowledge_length} tokens")
|
||
else:
|
||
# 如果不足长度，用空token填充
|
||
original_length = len(sentence_tokens)
|
||
sentence_tokens.extend([pad_token_id] * (knowledge_length - len(sentence_tokens)))
|
||
if original_length < knowledge_length:
|
||
Logger(f"Sentence {i+1} padded from {original_length} to {knowledge_length} tokens")
|
||
|
||
processed_rows.append(sentence_tokens)
|
||
|
||
if (i + 1) % 1000 == 0:
|
||
Logger(f"Processed {i + 1}/{num_to_process} sentences")
|
||
|
||
# 如果句子数量不足，用空token填充剩余位置
|
||
while len(processed_rows) < knowledge_num:
|
||
empty_tokens = [pad_token_id] * knowledge_length
|
||
processed_rows.append(empty_tokens)
|
||
if len(processed_rows) % 1000 == 0:
|
||
Logger(f"Added empty entry {len(processed_rows)}/{knowledge_num}")
|
||
|
||
Logger(f"Finished adding empty entries. Total: {len(processed_rows)}/{knowledge_num}")
|
||
|
||
# 转换为tensor
|
||
processed_tensor = torch.tensor(processed_rows, dtype=torch.long)
|
||
|
||
Logger(f"Data processing completed:")
|
||
Logger(f" - Processed {num_to_process} sentences")
|
||
Logger(f" - Added {knowledge_num - num_to_process} empty entries")
|
||
Logger(f" - Final shape: {processed_tensor.shape}")
|
||
Logger(f" - Expected shape: ({knowledge_num}, {knowledge_length})")
|
||
|
||
# 保存处理结果到缓存文件
|
||
try:
|
||
torch.save(processed_tensor, args.cluster_cache_path)
|
||
Logger(f"Processed results saved to {args.cluster_cache_path}")
|
||
except Exception as e:
|
||
Logger(f"Failed to save processed results: {e}")
|
||
|
||
# 4. 初始化模型的knowledge_dataset
|
||
if hasattr(model, 'knowledge_dataset') and hasattr(model.knowledge_dataset, 'knowledge_dataset'):
|
||
model.knowledge_dataset.knowledge_dataset.data.copy_(processed_tensor)
|
||
Logger("Successfully initialized model.knowledge_dataset.knowledge_dataset with processed data")
|
||
else:
|
||
Logger("Warning: Could not find model.knowledge_dataset.knowledge_dataset to initialize")
|
||
# 存储为全局变量作为备选
|
||
globals()['processed_database'] = processed_tensor
|
||
|
||
Logger(f"Database embeddings and sentences stored in model")
|
||
|
||
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, swanlab_run):
|
||
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 ''
|
||
best_loss = float('10000')
|
||
|
||
# 初始化CUDA事件变量
|
||
data_start = data_end = forward_start = forward_end = None
|
||
backward_start = backward_end = optimizer_start = optimizer_end = None
|
||
|
||
# 添加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 = []
|
||
|
||
# 记录初始内存状态
|
||
if args.memory_monitor:
|
||
log_memory_status(-1, prefetch_batches, accelerator, "before_prefetch", detailed=True)
|
||
|
||
# 预取初始批次
|
||
for i in range(min(prefetch_factor, len(train_loader))):
|
||
try:
|
||
batch = next(data_iter)
|
||
prefetch_batches.append(batch)
|
||
# 每次添加batch后记录内存变化
|
||
if args.memory_monitor and accelerator.is_main_process:
|
||
log_memory_status(-1, prefetch_batches, accelerator, f"after_adding_batch_{i+1}")
|
||
except StopIteration:
|
||
break
|
||
|
||
# 记录预取完成后的内存状态
|
||
if args.memory_monitor:
|
||
log_memory_status(-1, prefetch_batches, accelerator, "after_initial_prefetch", detailed=True)
|
||
|
||
# 在开始循环前初始化日志记录所需变量
|
||
last_log_time = epoch_start_time
|
||
|
||
for step in range(total_steps_in_epoch):
|
||
try:
|
||
# 计时数据加载 (只在主进程进行)
|
||
if args.profile and accelerator.is_main_process and data_start is not None:
|
||
data_start.record()
|
||
|
||
# 记录使用batch前的内存状态（根据配置间隔记录详细信息）
|
||
if args.memory_monitor and step % args.memory_monitor_interval == 0:
|
||
log_memory_status(step, prefetch_batches, accelerator, "before_use_batch", detailed=True)
|
||
|
||
# 使用预取的数据
|
||
if prefetch_batches:
|
||
X, Y, loss_mask = prefetch_batches.pop(0)
|
||
# 记录使用batch后的内存变化
|
||
if args.memory_monitor and step % args.memory_monitor_interval == 0:
|
||
log_memory_status(step, prefetch_batches, accelerator, "after_pop_batch")
|
||
else:
|
||
# 如果预取队列为空，直接加载
|
||
X, Y, loss_mask = next(data_iter)
|
||
if args.memory_monitor and accelerator.is_main_process:
|
||
Logger(f"[Memory Monitor] Step {step} - Prefetch queue empty, loading directly!", accelerator)
|
||
|
||
# 异步预取下一批数据
|
||
if step + prefetch_factor < len(train_loader):
|
||
try:
|
||
batch = next(data_iter)
|
||
prefetch_batches.append(batch)
|
||
# 记录添加新batch后的内存变化
|
||
if args.memory_monitor and step % args.memory_monitor_interval == 0:
|
||
log_memory_status(step, prefetch_batches, accelerator, "after_add_batch")
|
||
except StopIteration:
|
||
pass
|
||
|
||
# 计时数据加载结束 (只在主进程进行)
|
||
if args.profile and accelerator.is_main_process and data_end is not None:
|
||
data_end.record()
|
||
|
||
# 更新学习率
|
||
if scheduler is not None:
|
||
scheduler.step()
|
||
|
||
# 计时前向传播 (只在主进程进行)
|
||
if args.profile and accelerator.is_main_process and forward_start is not None:
|
||
forward_start.record()
|
||
|
||
# 前向传播
|
||
with ctx:
|
||
if step == 0 and args.embedding_epoch == epoch:
|
||
# 需要设置原始模型的freeze_embedding属性，而不是包装后的模型
|
||
unwrapped_model = accelerator.unwrap_model(model)
|
||
unwrapped_model.freeze_embedding = True
|
||
Logger(f"Set freeze_embedding=True for epoch {epoch}, step {step}", accelerator)
|
||
res = model(X, step=step)
|
||
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 and forward_end is not None:
|
||
forward_end.record()
|
||
|
||
# 计时反向传播 (只在主进程进行)
|
||
if args.profile and accelerator.is_main_process and backward_start is not None:
|
||
backward_start.record()
|
||
|
||
# 反向传播
|
||
# 当使用DeepSpeed时，它会自动处理梯度累积和梯度裁剪
|
||
accelerator.backward(loss)
|
||
|
||
# 计时反向传播结束 (只在主进程进行)
|
||
if args.profile and accelerator.is_main_process and backward_end is not None:
|
||
backward_end.record()
|
||
|
||
# 计时优化器步骤 (只在主进程进行)
|
||
if args.profile and accelerator.is_main_process and optimizer_start is not None:
|
||
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 and optimizer_end is not None:
|
||
optimizer_end.record()
|
||
|
||
# 打印训练信息 (只在主进程进行)
|
||
if (step + 1) % args.log_interval == 0 and accelerator.is_main_process:
|
||
current_time = time.time()
|
||
|
||
# 记录日志输出时的详细内存状态
|
||
if args.memory_monitor:
|
||
log_memory_status(step, prefetch_batches, accelerator, "at_log_interval", detailed=True)
|
||
|
||
# 强制垃圾回收并记录内存变化
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
gc.collect()
|
||
log_memory_status(step, prefetch_batches, accelerator, "after_gc", detailed=True)
|
||
|
||
# 计算性能指标
|
||
if args.profile and accelerator.is_main_process:
|
||
torch.cuda.synchronize()
|
||
|
||
# 确保所有事件都已记录才计算elapsed_time
|
||
try:
|
||
data_time = data_start.elapsed_time(data_end) if data_start is not None and data_end is not None else 0
|
||
forward_time = forward_start.elapsed_time(forward_end) if forward_start is not None and forward_end is not None else 0
|
||
backward_time = backward_start.elapsed_time(backward_end) if backward_start is not None and backward_end is not None else 0
|
||
optimizer_time = optimizer_start.elapsed_time(optimizer_end) if optimizer_start is not None and optimizer_end is not None else 0
|
||
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)
|
||
except RuntimeError as e:
|
||
if "Both events must be recorded" in str(e):
|
||
Logger(f"Warning: CUDA events not properly recorded, skipping performance analysis: {e}", accelerator)
|
||
else:
|
||
raise e
|
||
|
||
|
||
# 计算当前学习率
|
||
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 # 更新上次日志时间
|
||
|
||
log_dict = {
|
||
"epoch": epoch + 1,
|
||
"step": step + 1,
|
||
"total_steps_in_epoch": total_steps_in_epoch,
|
||
"loss": loss.item() * args.accumulation_steps,
|
||
"lr": current_lr,
|
||
"tokens_per_sec": tokens_per_sec,
|
||
"epoch_time_left_seconds": epoch_remaining_time,
|
||
"total_time_left_seconds": total_remaining_time
|
||
}
|
||
|
||
Logger(f"Epoch {epoch+1}/{args.epochs}, Step {step+1}/{total_steps_in_epoch}, "
|
||
f"Loss: {log_dict['loss']:.4f}, "
|
||
f"LR: {log_dict['lr']:.6f}, "
|
||
f"Speed: {log_dict['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 args.use_swanlab and accelerator.is_main_process and swanlab_run:
|
||
swanlab_run.log(log_dict)
|
||
|
||
# 保存模型 (只在主进程进行)
|
||
loss_total = loss.item() * args.accumulation_steps
|
||
if epoch > 1 and best_loss > loss_total and accelerator.is_main_process:
|
||
best_loss = loss_total
|
||
# 使用函数开始处定义的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)
|
||
# 记录异常时的内存状态
|
||
if args.memory_monitor:
|
||
log_memory_status(step, prefetch_batches, accelerator, "at_exception", detailed=True)
|
||
import traceback
|
||
Logger(traceback.format_exc(), accelerator)
|
||
|
||
# 清理prefetch_batches，防止内存泄漏
|
||
if args.memory_monitor and accelerator.is_main_process:
|
||
Logger(f"[Memory Monitor] Clearing prefetch_batches due to exception. Current length: {len(prefetch_batches)}", accelerator)
|
||
prefetch_batches.clear()
|
||
gc.collect()
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
if args.memory_monitor:
|
||
log_memory_status(step, prefetch_batches, accelerator, "after_exception_cleanup", detailed=True)
|
||
|
||
# 训练epoch结束时清理prefetch_batches
|
||
if args.memory_monitor:
|
||
if accelerator.is_main_process:
|
||
Logger(f"[Memory Monitor] Epoch {epoch+1} finished. Clearing prefetch_batches. Final length: {len(prefetch_batches)}", accelerator)
|
||
log_memory_status(total_steps_in_epoch-1, prefetch_batches, accelerator, "before_epoch_end_cleanup", detailed=True)
|
||
prefetch_batches.clear()
|
||
gc.collect()
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
if args.memory_monitor:
|
||
log_memory_status(total_steps_in_epoch-1, prefetch_batches, accelerator, "after_epoch_end_cleanup", detailed=True)
|
||
|
||
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=4)
|
||
parser.add_argument("--embedding_epoch", type=int, default=2, help="embedding训练的epoch数")
|
||
parser.add_argument("--batch_size", type=int, default=64)
|
||
parser.add_argument("--learning_rate", type=float, default=2e-4)
|
||
parser.add_argument("--dtype", type=str, default="bfloat16")
|
||
parser.add_argument("--use_swanlab", default=True, action="store_true") # 替换wandb参数
|
||
parser.add_argument("--swanlab_project", type=str, default="MiniMind-Pretrain") # 替换wandb参数
|
||
parser.add_argument("--num_workers", type=int, default=8)
|
||
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=512, type=int)
|
||
parser.add_argument('--n_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('--disable_db', action='store_true', help="禁用数据库功能，使用固定值1e-4替代")
|
||
parser.add_argument("--data_path", type=str, default="./dataset/merged_pretrain.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("--knowledge_num", type=int, default=960400,help="知识库的数据数目")
|
||
parser.add_argument("--knowledge_length", type=int, default=32,help="知识库的句子长度")
|
||
parser.add_argument("--database_init_path", type=str, default="./dataset/combined_prepare.json", help="数据库初始化路径")
|
||
parser.add_argument("--fast_clustering", action="store_true", default=True, help="使用快速近似聚类算法（适用于大数据集）")
|
||
parser.add_argument("--cluster_cache_path", type=str, default="./cache/cluster_tokens_single.pt", help="聚类结果缓存文件路径")
|
||
parser.add_argument("--recompute_clusters", action="store_true", default=False, help="强制重新计算聚类，忽略缓存文件")
|
||
parser.add_argument("--memory_monitor", action="store_true", default=False, help="启用内存监控")
|
||
parser.add_argument("--memory_monitor_interval", type=int, default=10, help="内存监控间隔（步数）")
|
||
args = parser.parse_args()
|
||
|
||
#########################################################
|
||
# 初始化accelerator和deepspeed
|
||
#########################################################
|
||
# 设置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="none", # 将优化器状态卸载到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,
|
||
knowledge_num=args.knowledge_num,
|
||
knowledge_length=args.knowledge_length,
|
||
embeddings_epoch=args.embedding_epoch
|
||
)
|
||
|
||
#########################################################
|
||
# 创建保存目录
|
||
#########################################################
|
||
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)
|
||
|
||
#########################################################
|
||
# 设置数据类型
|
||
#########################################################
|
||
pt_dtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[args.dtype]
|
||
|
||
|
||
#########################################################
|
||
# 配置SwanLab
|
||
#########################################################
|
||
# 设置SwanLab运行名称
|
||
args.swanlab_run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
|
||
|
||
# 合并args和lm_config为一个字典（无论是否使用SwanLab都需要，用于打印配置信息）
|
||
config_dict = vars(args).copy()
|
||
config_dict.update(vars(lm_config))
|
||
|
||
# 初始化SwanLab实验实例
|
||
swanlab_run = None
|
||
if args.use_swanlab and accelerator.is_main_process:
|
||
# 初始化SwanLab
|
||
swanlab_run = swanlab.init(
|
||
project=args.swanlab_project,
|
||
experiment_name=args.swanlab_run_name,
|
||
description="MiniMind预训练实验，使用本地部署的SwanLab进行可视化",
|
||
config=config_dict
|
||
# 设置SwanLab服务器地址和API Key
|
||
# host="http://100.123.118.114:11071",
|
||
# api_key="LesBT7HRq23HNBrOPKP8S"
|
||
)
|
||
else:
|
||
swanlab_run = None
|
||
|
||
#########################################################
|
||
# 打印信息
|
||
#########################################################
|
||
# 计算每次迭代的token数量
|
||
tokens_per_iter = args.batch_size * lm_config.max_seq_len
|
||
if accelerator.is_main_process:
|
||
Logger(f"tokens_per_iter: {tokens_per_iter}", accelerator)
|
||
Logger("Configuration:", accelerator)
|
||
for key, value in config_dict.items():
|
||
Logger(f" {key}: {value}", accelerator)
|
||
|
||
|
||
#########################################################
|
||
# 设置自动混合精度上下文
|
||
#########################################################
|
||
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, args.database_init_path, args)
|
||
# 将accelerator传递给init_model函数中的Logger调用
|
||
Logger(f'模型初始化完成', accelerator)
|
||
|
||
#########################################################
|
||
# 处理位置编码张量问题
|
||
#########################################################
|
||
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
|
||
)
|
||
|
||
#########################################################
|
||
# 训练循环
|
||
#########################################################
|
||
overall_start_time = time.time() # Record overall start time
|
||
for epoch in range(args.epochs):
|
||
Logger(f"开始第{epoch+1}轮训练", accelerator)
|
||
train_epoch(epoch, accelerator, model, train_loader, optimizer, scheduler, args, ctx, overall_start_time, swanlab_run) # Pass overall start time
|
||
|
||
# 每个epoch结束后进行内存清理
|
||
Logger(f"第{epoch+1}轮训练完成，进行内存清理", accelerator)
|
||
gc.collect()
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
|
||
# 记录epoch结束时的内存状态
|
||
if accelerator.is_main_process:
|
||
memory_info = get_memory_usage()
|
||
cuda_info = get_cuda_memory_usage()
|
||
log_msg = f"[Memory Monitor] Epoch {epoch+1} completed - "
|
||
log_msg += f"System RSS: {memory_info['rss_mb']:.2f}MB"
|
||
if cuda_info:
|
||
log_msg += f", CUDA allocated: {cuda_info['cuda_allocated_mb']:.2f}MB"
|
||
log_msg += f", CUDA reserved: {cuda_info['cuda_reserved_mb']:.2f}MB"
|
||
Logger(log_msg, accelerator)
|
||
|
||
#########################################################
|
||
# 关闭SwanLab
|
||
#########################################################
|
||
if args.use_swanlab and accelerator.is_main_process and swanlab_run:
|
||
swanlab_run.finish()
|
||
|
||
if __name__ == "__main__":
|
||
main()
|