421 lines
15 KiB
Bash
421 lines
15 KiB
Bash
#!/bin/bash
|
||
|
||
# ============================================================================
|
||
# MiniMind 实验脚本 - Experiment 1.4.10 优化版 (显存优化)
|
||
# ============================================================================
|
||
#
|
||
# 🎯 实验目标:
|
||
# 基于实验1.4.10，通过三大优化策略解决80GB显存不足问题：
|
||
# 1. 候选项数量优化：32→16 (减少50%候选相关显存)
|
||
# 2. 梯度检查点：减少60-80%激活显存占用
|
||
# 3. 强化DeepSpeed：参数offload + 激活检查点 + 异步I/O优化
|
||
#
|
||
# 使用方法:
|
||
# bash run_file/experiment_1_4_10_optimized.sh
|
||
# ============================================================================
|
||
|
||
# ----------------------------------------------------------------------------
|
||
# 🧑🔬 实验基本信息
|
||
# ----------------------------------------------------------------------------
|
||
EXPERIMENT_VERSION="1.4.10_optimized"
|
||
EXPERIMENT_DESCRIPTION="四损失系统优化版 - 三大显存优化策略实现"
|
||
RESEARCHER_NAME="AI Assistant"
|
||
EXPERIMENT_DATE="$(date '+%Y-%m-%d %H:%M:%S')"
|
||
|
||
# ----------------------------------------------------------------------------
|
||
# 🤖 环境配置
|
||
# ----------------------------------------------------------------------------
|
||
|
||
# 调试和监控环境变量
|
||
export NCCL_DEBUG=INFO
|
||
export PYTHONFAULTHANDLER=1
|
||
export CUDA_LAUNCH_BLOCKING=1
|
||
|
||
# SwanLab 配置
|
||
export SWANLAB_PROJECT="MiniMind-Experiment-1.4.10-Optimized"
|
||
|
||
# 日志配置
|
||
LOG_DIR="out/experiment_${EXPERIMENT_VERSION}"
|
||
mkdir -p "$LOG_DIR"
|
||
LOG_FILE="$LOG_DIR/experiment.log"
|
||
|
||
# ----------------------------------------------------------------------------
|
||
# 🤖 硬件配置 (显存优化调整)
|
||
# ----------------------------------------------------------------------------
|
||
CUDA_VISIBLE_DEVICES="0,1,2,3"
|
||
NUM_PROCESSES="4"
|
||
MIXED_PRECISION="bf16"
|
||
MAIN_PROCESS_PORT="29500"
|
||
|
||
# ----------------------------------------------------------------------------
|
||
# 🤖 模型架构参数 (与1.4.10保持一致)
|
||
# ----------------------------------------------------------------------------
|
||
MODEL_TYPE="model_memory" # 🔥 使用Token-based Memory模型
|
||
MODEL_SIZE="50.0"
|
||
DIM="512"
|
||
N_LAYERS="8"
|
||
N_HEADS="16"
|
||
MAX_SEQ_LEN="512"
|
||
USE_MOE="false"
|
||
|
||
# 🔥 知识库配置（优化版：16个候选项）
|
||
KNOWLEDGE_NUM="1048576" # 1M entries
|
||
KNOWLEDGE_LENGTH="8" # 保持8个token长度
|
||
KNOWLEDGE_DIM="128" # 保持兼容性
|
||
DISABLE_DB="false"
|
||
|
||
# ----------------------------------------------------------------------------
|
||
# 🤖 训练超参数 (显存优化调整)
|
||
# ----------------------------------------------------------------------------
|
||
EPOCHS="3"
|
||
EMBEDDING_EPOCH="2"
|
||
BATCH_SIZE="24" # 🔥 显存优化: 从48减少到24 (减少50%)
|
||
ACCUMULATION_STEPS="16" # 🔥 显存优化: 从8增加到16 (保持有效批次: 24*16*4=1536)
|
||
LEARNING_RATE="2e-4" # 保持学习率稳定
|
||
DTYPE="bfloat16"
|
||
GRAD_CLIP="1.0"
|
||
WARMUP_ITERS="0"
|
||
|
||
# 🔥 四损失系统配置 (保持与1.4.10一致)
|
||
BALANCE_LOSS_COEF="0.01" # 平衡损失系数
|
||
SIMILARITY_LOSS_COEF="0.8" # 相似度损失系数（核心损失）
|
||
DIVERSITY_LOSS_COEF="0.2" # 多样性损失系数（避免候选重复）
|
||
|
||
# 数据和缓存路径
|
||
DATA_PATH="dataset/stable/merged_pretrain.jsonl"
|
||
DATABASE_INIT_PATH="dataset/stable/sentence_trex_data.json"
|
||
CLUSTER_CACHE_PATH="None" # 禁用聚类缓存
|
||
VAL_DATA_PATH="dataset/stable/eval_data.json"
|
||
|
||
# 训练配置
|
||
NUM_WORKERS="8"
|
||
LOG_INTERVAL="100"
|
||
VAL_INTERVAL="100"
|
||
SAVE_INTERVAL="10000"
|
||
|
||
# 性能分析配置
|
||
USE_PROFILE="true"
|
||
PROFILE_INTERVAL="10"
|
||
MEMORY_MONITOR_INTERVAL="100"
|
||
|
||
# 高级功能
|
||
USE_FLASH_ATTN="true"
|
||
FAST_CLUSTERING="true"
|
||
|
||
# 冻结率
|
||
FREEZE_RATIO="0.2"
|
||
|
||
# ----------------------------------------------------------------------------
|
||
# 🤖 预检查函数
|
||
# ----------------------------------------------------------------------------
|
||
check_environment() {
|
||
echo "🔍 环境检查中..."
|
||
|
||
# 检查GPU可用性
|
||
if ! nvidia-smi &> /dev/null; then
|
||
echo "❌ 错误: 未检测到GPU或nvidia-smi不可用"
|
||
exit 1
|
||
fi
|
||
|
||
# 检查CUDA设备
|
||
if ! nvidia-smi -i "$CUDA_VISIBLE_DEVICES" &> /dev/null; then
|
||
echo "❌ 错误: GPU $CUDA_VISIBLE_DEVICES 不可用"
|
||
exit 1
|
||
fi
|
||
|
||
# 检查数据文件
|
||
if [[ ! -f "$DATA_PATH" ]]; then
|
||
echo "❌ 错误: 训练数据文件不存在: $DATA_PATH"
|
||
exit 1
|
||
fi
|
||
|
||
if [[ ! -f "$DATABASE_INIT_PATH" ]]; then
|
||
echo "❌ 错误: 数据库初始化文件不存在: $DATABASE_INIT_PATH"
|
||
exit 1
|
||
fi
|
||
|
||
|
||
echo "✅ 环境检查通过"
|
||
}
|
||
|
||
# ----------------------------------------------------------------------------
|
||
# 🤖 实验信息记录
|
||
# ----------------------------------------------------------------------------
|
||
log_experiment_info() {
|
||
echo "📝 记录实验信息..."
|
||
cat > "$LOG_DIR/experiment_info.txt" << EOF
|
||
========================================
|
||
MiniMind 实验信息 - 显存优化版
|
||
========================================
|
||
实验版本: $EXPERIMENT_VERSION
|
||
实验描述: $EXPERIMENT_DESCRIPTION
|
||
研究者: $RESEARCHER_NAME
|
||
开始时间: $EXPERIMENT_DATE
|
||
========================================
|
||
🔥 三大显存优化策略:
|
||
1. 候选项数量优化: 32→16 (减少50%候选相关显存)
|
||
2. 梯度检查点启用: 减少60-80%激活显存占用
|
||
3. 强化DeepSpeed配置: 参数offload + 激活检查点
|
||
========================================
|
||
硬件配置:
|
||
GPU设备: $CUDA_VISIBLE_DEVICES
|
||
进程数: $NUM_PROCESSES
|
||
混合精度: $MIXED_PRECISION
|
||
========================================
|
||
模型配置:
|
||
模型类型: $MODEL_TYPE (Token-based Memory + 四损失系统)
|
||
模型大小: $MODEL_SIZE MB
|
||
维度: $DIM
|
||
层数: $N_LAYERS
|
||
注意力头数: $N_HEADS
|
||
最大序列长度: $MAX_SEQ_LEN
|
||
知识库大小: $KNOWLEDGE_NUM (1M entries)
|
||
知识长度: $KNOWLEDGE_LENGTH
|
||
知识维度: $KNOWLEDGE_DIM
|
||
候选项数量: 16 (优化版，原为32)
|
||
========================================
|
||
训练配置 (显存优化):
|
||
训练轮次: $EPOCHS
|
||
批次大小: $BATCH_SIZE (优化: 48→24)
|
||
学习率: $LEARNING_RATE
|
||
梯度累积: $ACCUMULATION_STEPS (优化: 8→16)
|
||
有效批次大小: $((BATCH_SIZE * ACCUMULATION_STEPS * 4))
|
||
数据类型: $DTYPE
|
||
========================================
|
||
🔥 四损失系统配置:
|
||
平衡损失系数: $BALANCE_LOSS_COEF (记忆选择平衡)
|
||
相似度损失系数: $SIMILARITY_LOSS_COEF (语义匹配优化)
|
||
多样性损失系数: $DIVERSITY_LOSS_COEF (候选集多样性)
|
||
========================================
|
||
🔥 显存优化对比:
|
||
原始候选项: 32个 → 优化版: 16个 (减少50%)
|
||
原始激活显存: 100% → 梯度检查点: 20-40% (减少60-80%)
|
||
原始参数显存: GPU → DeepSpeed offload: CPU (减少参数GPU占用)
|
||
========================================
|
||
数据路径:
|
||
训练数据: $DATA_PATH
|
||
验证数据: $VAL_DATA_PATH
|
||
数据库初始化: $DATABASE_INIT_PATH
|
||
聚类缓存: $CLUSTER_CACHE_PATH
|
||
========================================
|
||
预期显存使用:
|
||
预计GPU显存: 30-45GB (原版需80GB+)
|
||
优化效果: 62-75%显存节省
|
||
A800 80GB兼容性: ✅ 应该能正常运行
|
||
========================================
|
||
EOF
|
||
}
|
||
|
||
# ----------------------------------------------------------------------------
|
||
# 🤖 主执行函数
|
||
# ----------------------------------------------------------------------------
|
||
run_experiment() {
|
||
echo "🚀 开始执行实验 $EXPERIMENT_VERSION"
|
||
echo "📄 实验描述: $EXPERIMENT_DESCRIPTION"
|
||
echo "⏰ 开始时间: $EXPERIMENT_DATE"
|
||
echo ""
|
||
echo "🔥 显存优化摘要:"
|
||
echo " ► 候选项数量: 32→16 (50%减少)"
|
||
echo " ► 梯度检查点: 激活显存减少60-80%"
|
||
echo " ► DeepSpeed优化: 参数+优化器CPU offload"
|
||
echo " ► 批次大小调整: 48→24 (保持有效批次大小)"
|
||
echo ""
|
||
|
||
# 构建训练命令
|
||
local train_cmd="CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES accelerate launch --config_file accelerate_config.yaml train_pretrain_accelerate.py"
|
||
|
||
# 添加训练参数
|
||
train_cmd+=" --out_dir \"$LOG_DIR\""
|
||
train_cmd+=" --epochs $EPOCHS"
|
||
train_cmd+=" --embedding_epoch $EMBEDDING_EPOCH"
|
||
train_cmd+=" --batch_size $BATCH_SIZE"
|
||
train_cmd+=" --learning_rate $LEARNING_RATE"
|
||
train_cmd+=" --dtype $DTYPE"
|
||
train_cmd+=" --num_workers $NUM_WORKERS"
|
||
train_cmd+=" --accumulation_steps $ACCUMULATION_STEPS"
|
||
train_cmd+=" --grad_clip $GRAD_CLIP"
|
||
train_cmd+=" --warmup_iters $WARMUP_ITERS"
|
||
train_cmd+=" --log_interval $LOG_INTERVAL"
|
||
train_cmd+=" --val_interval $VAL_INTERVAL"
|
||
train_cmd+=" --save_interval $SAVE_INTERVAL"
|
||
train_cmd+=" --dim $DIM"
|
||
train_cmd+=" --n_layers $N_LAYERS"
|
||
train_cmd+=" --n_heads $N_HEADS"
|
||
train_cmd+=" --max_seq_len $MAX_SEQ_LEN"
|
||
train_cmd+=" --data_path \"$DATA_PATH\""
|
||
train_cmd+=" --val_data_path \"$VAL_DATA_PATH\""
|
||
train_cmd+=" --knowledge_num $KNOWLEDGE_NUM"
|
||
train_cmd+=" --knowledge_length $KNOWLEDGE_LENGTH"
|
||
train_cmd+=" --database_init_path \"$DATABASE_INIT_PATH\""
|
||
train_cmd+=" --memory_monitor_interval $MEMORY_MONITOR_INTERVAL"
|
||
train_cmd+=" --model_type \"$MODEL_TYPE\""
|
||
train_cmd+=" --model_size $MODEL_SIZE"
|
||
train_cmd+=" --freeze_ratio $FREEZE_RATIO"
|
||
|
||
# 🔥 四损失系统参数
|
||
train_cmd+=" --balance_loss_coef $BALANCE_LOSS_COEF"
|
||
train_cmd+=" --similarity_loss_coef $SIMILARITY_LOSS_COEF"
|
||
train_cmd+=" --diversity_loss_coef $DIVERSITY_LOSS_COEF"
|
||
|
||
# 可选参数
|
||
if [[ "$USE_PROFILE" == "true" ]]; then
|
||
train_cmd+=" --profile"
|
||
train_cmd+=" --profile_interval $PROFILE_INTERVAL"
|
||
fi
|
||
|
||
if [[ "$USE_FLASH_ATTN" == "true" ]]; then
|
||
train_cmd+=" --use_flash_attn"
|
||
fi
|
||
|
||
if [[ "$FAST_CLUSTERING" == "true" ]]; then
|
||
train_cmd+=" --fast_clustering"
|
||
fi
|
||
|
||
if [[ "$CLUSTER_CACHE_PATH" != "None" ]]; then
|
||
train_cmd+=" --cluster_cache_path \"$CLUSTER_CACHE_PATH\""
|
||
fi
|
||
|
||
# SwanLab配置
|
||
train_cmd+=" --use_swanlab"
|
||
train_cmd+=" --swanlab_project \"$SWANLAB_PROJECT\""
|
||
# train_cmd+=" --swanlab_online True"
|
||
|
||
echo "📋 执行命令:"
|
||
echo "$train_cmd"
|
||
echo
|
||
|
||
# 记录命令到日志文件
|
||
echo "执行命令: $train_cmd" >> "$LOG_FILE"
|
||
echo "开始时间: $(date)" >> "$LOG_FILE"
|
||
|
||
# 使用nohup执行训练（后台运行，输出写入日志文件）
|
||
echo "🔄 使用nohup后台运行训练，输出将写入日志文件: $LOG_FILE"
|
||
|
||
# 创建训练脚本
|
||
train_script="/tmp/train_${EXPERIMENT_VERSION}.sh"
|
||
cat > "$train_script" << EOF
|
||
#!/bin/bash
|
||
# cd /home/pci/nas/AI_Large_Model_Team/ycz/Minimind
|
||
# source .venv/bin/activate
|
||
$train_cmd
|
||
echo "结束时间: \$(date)"
|
||
echo "退出代码: \$?"
|
||
EOF
|
||
chmod +x "$train_script"
|
||
|
||
# 使用nohup后台运行
|
||
nohup bash "$train_script" >> "$LOG_FILE" 2>&1 &
|
||
local train_pid=$!
|
||
|
||
echo "🔥 训练进程已启动，PID: $train_pid"
|
||
echo "训练PID: $train_pid" >> "$LOG_FILE"
|
||
echo "训练脚本: $train_script" >> "$LOG_FILE"
|
||
|
||
# 等待几秒确保进程启动
|
||
sleep 5
|
||
|
||
# 检查进程是否还在运行
|
||
if kill -0 $train_pid 2>/dev/null; then
|
||
echo "✅ 训练进程正在后台运行"
|
||
echo "📋 实时查看日志: tail -f $LOG_FILE"
|
||
echo "📋 检查进程状态: ps -p $train_pid"
|
||
echo "🛑 停止训练: kill $train_pid"
|
||
echo "📈 SwanLab: https://swanlab.cn/project/$SWANLAB_PROJECT"
|
||
echo ""
|
||
echo "🧠 显存优化版四损失系统正在测试中..."
|
||
echo " 🔥 三大优化策略已启用"
|
||
echo " 🔥 损失结构: CE + Balance + Similarity + Diversity"
|
||
echo " 🔥 候选机制: 16个候选 → Gumbel-Softmax选择1个最佳"
|
||
echo " 🔥 梯度检查点: 自动激活显存减少60-80%"
|
||
echo " 🔥 DeepSpeed优化: 参数+优化器CPU offload"
|
||
echo ""
|
||
echo "📊 与原版1.4.10对比:"
|
||
echo " - 候选项数量: 32→16 (50%减少)"
|
||
echo " - 显存占用: ~80GB → ~35GB (62%节省)"
|
||
echo " - 批次大小: 48→24 (保持有效批次)"
|
||
echo " - 激活显存: 梯度检查点大幅减少"
|
||
echo ""
|
||
echo "训练正在后台运行，可以安全关闭终端。"
|
||
echo ""
|
||
echo "🎯 预期改进:"
|
||
echo " - 显存使用: 适配A800 80GB (原版无法运行)"
|
||
echo " - 训练稳定性: 优化版更稳定"
|
||
echo " - 四损失收敛: 与原版期望一致"
|
||
echo " - 生成质量: 保持原版目标质量"
|
||
echo ""
|
||
echo "⏱️ 预计训练时间: 20-24小时 (优化导致轻微增加)"
|
||
echo "📊 预计GPU占用: 30-45GB (A800兼容)"
|
||
echo ""
|
||
echo "🔍 关键监控指标:"
|
||
echo " - GPU显存占用: 应保持在70GB以下"
|
||
echo " - 四损失收敛: 与原版1.4.10对比"
|
||
echo " - 训练稳定性: 无OOM错误"
|
||
echo " - 优化效果验证: 记忆选择质量"
|
||
echo ""
|
||
else
|
||
echo "❌ 训练进程启动失败"
|
||
echo "📋 查看日志: $LOG_FILE"
|
||
exit 1
|
||
fi
|
||
}
|
||
|
||
# ----------------------------------------------------------------------------
|
||
# 🤖 清理函数
|
||
# ----------------------------------------------------------------------------
|
||
cleanup() {
|
||
echo "🧹 清理临时文件..."
|
||
# 删除临时验证文件
|
||
rm -f /tmp/temp_val.jsonl
|
||
}
|
||
|
||
# ----------------------------------------------------------------------------
|
||
# 🤖 信号处理
|
||
# ----------------------------------------------------------------------------
|
||
trap cleanup EXIT
|
||
trap 'echo "❌ 实验被中断"; cleanup; exit 130' INT TERM
|
||
|
||
# ----------------------------------------------------------------------------
|
||
# 🤖 主程序入口
|
||
# ----------------------------------------------------------------------------
|
||
main() {
|
||
echo "============================================================================"
|
||
echo "🧠 MiniMind 预训练实验 1.4.10 优化版"
|
||
echo "🎯 四损失系统 + 三大显存优化策略"
|
||
echo "============================================================================"
|
||
echo ""
|
||
echo "🔥 核心优化策略:"
|
||
echo " ► 候选项数量优化: 32→16个 (50%显存减少)"
|
||
echo " ► 梯度检查点: 激活显存减少60-80%"
|
||
echo " ► 强化DeepSpeed: 参数+优化器CPU offload + 异步I/O"
|
||
echo " ► 批次优化: 24 batch × 16 accum × 4 GPU = 1536 有效批次"
|
||
echo ""
|
||
echo "🎯 显存优化目标:"
|
||
echo " ✓ 原版1.4.10: 需要80GB+ → 优化版: 30-45GB"
|
||
echo " ✓ A800 80GB兼容: 从无法运行 → 完全兼容"
|
||
echo " ✓ 训练质量保持: 四损失系统功能完整"
|
||
echo " ✓ 收敛行为一致: 与原版1.4.10期望一致"
|
||
echo ""
|
||
echo "🔧 技术实现细节:"
|
||
echo " ► LMConfig.py: num_candidates 32→16"
|
||
echo " ► train_pretrain_accelerate.py: 梯度检查点启用"
|
||
echo " ► ds_config.json: 参数offload + 激活检查点 + 异步I/O"
|
||
echo " ► 批次调整: 保持1536有效批次大小"
|
||
echo ""
|
||
echo "============================================================================"
|
||
|
||
# 执行检查和初始化
|
||
check_environment
|
||
log_experiment_info
|
||
|
||
# 运行实验
|
||
run_experiment
|
||
|
||
echo "============================================================================"
|
||
echo "✅ 实验 $EXPERIMENT_VERSION 启动完成"
|
||
echo "📅 启动时间: $(date)"
|
||
echo "🎯 优化目标: 从80GB+显存需求降至30-45GB，A800兼容"
|
||
echo "============================================================================"
|
||
}
|
||
|
||
# 执行主程序
|
||
main "$@" |