Minimind/CLAUDE.md

CLAUDE.md - MiniMind 预训练项目指南

项目概述: MiniMind 大语言模型预训练项目，研究使用人类可理解的 KnowledgeDataset 替代传统 Transformer Feed-Forward 层作为记忆层。

📋 目录

• 项目架构
• 环境配置
• 训练流程
• 实验管理
• 配置参数
• 故障排除

🏗️ 项目架构

核心模型

文件	用途	说明
model/model.py	主要模型	Transformer + KnowledgeDataset 记忆层
model/model_no_feed.py	无FFN变体	不使用 Feed-Forward 层的实验版本
model/model_original.py	基线模型	传统 Transformer 架构（实验对照）
model/LMConfig.py	配置管理	支持 MOE、数据库、知识图谱功能
model/dataset.py	数据处理	预训练数据集加载和处理

关键特性

• ✨ 人类可理解记忆层: 使用 KnowledgeDataset 替代传统 FFN
• 🚀 分布式训练: Accelerate + DeepSpeed 支持
• 📊 实时监控: SwanLab 训练可视化
• 🔧 灵活配置: 支持多种模型架构实验

目录结构

pretrains-worktree/
├── model/                    # 模型定义
│   ├── model.py             # 主要模型（含KnowledgeDataset）
│   ├── model_original.py    # 基线模型
│   ├── model_no_feed.py     # 无FFN变体
│   ├── LMConfig.py          # 配置类
│   └── dataset.py           # 数据集处理
├── preprocessing/           # 数据预处理
├── run_file/               # 实验脚本
├── out/                    # 输出目录
├── accelerate_config.yaml  # 分布式配置
├── ds_config.json         # DeepSpeed配置
├── train_pretrain_accelerate.py  # 主训练脚本
└── eval_model.py           # 模型推理评估脚本

🔬 研究现状

研究重点

• KnowledgeDataset: 探索人类可理解的神经网络记忆机制

当前问题

1. 文本生成质量:

   • Loss 收敛良好 (model: 0.6 vs baseline: 1.9)
   • 但输出文本为词组碎片，缺乏句法连贯性

2. SFT 效果差异:

   • model 的 SFT 效果远低于 model_original 基线

⚙️ 环境配置

1. 环境管理

# 使用 uv 包管理器的 .venv 环境

# 添加新包
uv add <package_name>

# 同步环境
uv sync

2. 数据预处理

# 预处理预训练数据
python preprocessing/preprocess_pretrain.py

# 预处理三元组数据
python preprocessing/preprocess_trex.py

# 预处理组合数据
python preprocessing/preprocess_combined_json.py

🚀 训练流程

快速开始

# 执行实验脚本
bash run_file/experiment_1.4.XX.sh

🧪 实验管理

核心文件

• 实验记录模版: experiment/EXPERIMENT_TEMPLATE.md - 标准化的实验记录格式
• 实验脚本模版: run_file/experiment_template.sh - 自动化的实验执行脚本
• 管理指南: experiment/README.md - 详细的实验管理流程说明

🤝 人类-AI 协作模式

🧑‍🔬 人类职责（最简化）

1. 填写实验目标 - 在实验记录中填写：
   • 基于实验（上一版实验编号）
   • 实验目的、研究假设、预期结果
2. 审核确认 - 审核AI生成的完整记录
3. 提交决策 - 决定是否git commit

🤖 AI职责（全流程管理）

1. 实验设计 - 记录详细的思考过程和决策逻辑
2. 脚本管理 - 完全负责生成和管理实验脚本
3. 执行监控 - 实时记录训练过程和资源使用
4. 结果分析 - 自动分析性能指标和问题诊断
5. Git记录 - 生成代码变更记录和版本对比

实验流程

# 1. 人类确定实验版本和目标
EXPERIMENT_VERSION="1.4.1"

# 2. AI创建实验文件
cp experiment/EXPERIMENT_TEMPLATE.md experiment/experiment_${EXPERIMENT_VERSION}.md
cp run_file/experiment_template.sh run_file/experiment_${EXPERIMENT_VERSION}.sh

# 3. 人类填写基本信息（仅需填写[人类填写]部分）

# 4. AI完成所有技术工作：
#    - 思考过程记录
#    - 参数配置
#    - 脚本生成
#    - 实验执行（使用nohup后台运行）
#    - 结果分析

# 5. 人类审核 -> AI提交git

🔧 后台训练执行

使用nohup确保训练持续进行

所有实验脚本现已集成nohup后台运行功能：

# 执行实验（自动使用nohup后台运行）
bash run_file/experiment_X.X.X.sh

# 实时监控训练进度
tail -f out/experiment_X_X_X/experiment.log

# 检查训练进程状态
ps aux | grep train_pretrain_accelerate

# 手动停止训练（如需要）
kill [PID]

重要特性

• ✅ 后台运行: 使用nohup确保训练在SSH断开后继续
• 📝 日志记录: 所有输出自动记录到实验日志文件
• 🔍 进程监控: 提供PID和状态检查命令
• 🛑 优雅停止: 支持安全的训练中断机制
• ⏰ 时间估算: 自动显示预计训练完成时间

实验记录结构

experiment_X.Y.Z.md
├── 🧠 AI思考过程        # AI的设计思路和决策推理
├── 📝 Git变更记录       # 代码修改详情和原因
├── 📋 实验基本信息      # 人类填写目标，AI填写配置
├── ⚙️ 配置参数         # AI根据目标自动配置
├── 🚀 执行记录         # 训练过程实时更新
├── 📊 训练结果         # 自动化的结果分析
├── 🔍 推理评估         # 使用eval_model.py的实际推理效果
├── 📈 深度分析         # 问题诊断和改进建议
└── 🎯 实验结论         # 假设验证和后续计划

🔍 实验评估要求

重要: 每个实验在训练完成后，必须运行 eval_model.py 进行实际推理效果评估：

# 基本评估命令（使用默认参数）
.venv/bin/python eval_model.py \
    --model_path out/experiment_X_Y_Z/pretrain_512.pth \
    --model_type model

# 完整评估命令（指定所有参数）
.venv/bin/python eval_model.py \
    --model_path out/experiment_X_Y_Z/pretrain_512.pth \
    --model_type model \
    --dim 512 \
    --n_layers 8 \
    --n_heads 32 \
    --knowledge_num 1048576 \
    --knowledge_length 32 \
    --knowledge_dim 128

评估指标说明

• 输入/输出对比: 展示模型对前30个token的续写能力
• Loss值: 量化预测准确度，越低越好
• 文本连贯性: 观察生成文本是否符合语法和语义
• 模型对比: 比较model、model_original、model_no_feed的差异

📋 模型版本管理

重要: 为了方便使用 eval_model.py 对不同版本的模型进行测试，每一次实验后需要把模型文件拷贝为 model_X_X_X.py 文件，这样以后就可以通过修改 eval_model.py 来重新使用旧的模型文件来推理。

模型文件拷贝流程

# 实验完成后，将当前模型文件拷贝为版本化文件
# 例如：实验1.4.4完成后
cp model/model.py model/model_1_4_4.py

# 或者如果使用了其他变体
cp model/model_memory.py model/model_memory_1_4_4.py

使用版本化模型进行评估

# 评估历史版本模型（需要先在eval_model.py中添加对应的import支持）
.venv/bin/python eval_model.py \
    --model_path out/experiment_1_4_4/pretrain_512.pth \
    --model_type model_1_4_4 \
    --num_samples 10

模型版本命名规范

原始文件	版本化文件名	用途说明
model/model.py	model/model_X_Y_Z.py	主要模型的历史版本
model/model_memory.py	model/model_memory_X_Y_Z.py	记忆模型的历史版本
model/model_original.py	model/model_original_X_Y_Z.py	基线模型的历史版本

注意:

• 使用版本化模型前，需要在 eval_model.py 中添加相应的 import 语句
• 版本号格式与实验版本号保持一致（如 1_4_4，使用下划线分隔）
• 建议在实验记录中注明使用的模型文件版本

版本命名规范

版本格式	说明	示例
X.Y.Z	主要.次要.修订	1.4.1
主要版本 (X)	重大架构变更	从 model_original 到 model
次要版本 (Y)	功能增强或重要参数调整	新增知识库功能
修订版本 (Z)	小幅调整和优化	学习率调整、批次大小优化

质量标准

✅ 合格实验必须满足:

• 明确的实验目标和可验证假设
• 完整的AI思考过程记录
• 详细的Git变更记录
• 训练过程稳定且结果可解释
• 运行eval_model.py进行推理评估
• 具体可行的改进建议

❌ 不合格情况:

• 目标模糊或无法验证
• 缺少思考过程或Git记录
• 训练异常中断或数据错误
• 未进行推理评估或缺少评估结果
• 结论不明确或缺乏下一步计划

⚙️ 配置参数

配置文件

文件	用途
accelerate_config.yaml	Accelerate 分布式训练配置
ds_config.json	DeepSpeed ZeRO Stage 2 优化配置
pyproject.toml	项目依赖和环境配置

硬件配置 (单张 RTX 4090)

核心参数

参数类别	参数名	值	说明
训练设置	epochs	3	训练轮次
	batch_size	128	批次大小
	accumulation_steps	8	梯度累积步数
	mixed_precision	bf16	混合精度训练
模型架构	dim	512	模型维度
	n_layers	8	Transformer 层数
	n_heads	≤32	注意力头数
	max_seq_len	512	最大序列长度
知识库	knowledge_num	1048576	知识条目数量
	knowledge_length	32	单条知识长度
其他	use_moe	false	不使用专家混合

数据路径

# 预训练数据
data_path="/home/pci/ycz/Code/Minimind/dataset/stable/merged_pretrain.jsonl"

# 知识库初始化
database_init_path="/home/pci/ycz/Code/Minimind/dataset/stable/sentence_trex_data.json"

# 聚类缓存（可选）
cluster_cache_path=None  # 默认关闭

📊 训练监控

SwanLab 可视化

• ✅ 训练指标: 实时监控 loss、学习率变化
• 📈 资源监控: GPU 内存、计算利用率追踪
• 🌐 多模式: 支持在线/离线监控模式

🛠️ 故障排除

⚠️ 记忆增强架构的关键设计原则

核心教训: 基于实验1.4.3灾难性失败的深刻反思

🎯 查询机制特异性原则

最重要的架构设计原则：在记忆增强的语言模型中，查询机制的特异性比融合机制的复杂性更加重要。

架构选择	查询输入	记忆选择特性	训练表现	推理表现	结果评价
✅ 正确	h_attn	多样化、上下文相关	健康收敛	良好泛化	可用架构
❌ 错误	x + h_attn	固化、选择单一	"完美"记忆化	灾难性失败	禁用架构

🚨 灾难性过拟合的识别与预防

早期预警信号：

危险信号	安全阈值	危险阈值	建议行动
训练Loss过低	>0.5	<0.1	⛔ 立即停止训练
训练-推理Loss差异	<5倍	>10倍	⛔ 回滚架构修改
生成文本重复率	<50%	>80%	⛔ 检查记忆选择固化
记忆选择熵值	>3.0	<2.0	⛔ 增加查询多样性

实验1.4.3的教训：

• 训练Loss: 0.006 (极度危险)
• 推理Loss: 29.34 (4890倍差异)
• 生成质量: 0/10 (完全失败)
• 根本原因: h = x + h_attn 导致查询向量平均化，记忆选择完全固化

🛡️ 架构设计防护规则

记忆查询输入选择：

# ✅ 推荐：使用注意力输出作为记忆查询
query = h_attn  # 保持内容相关性和位置特异性

# ❌ 禁止：使用混合信息作为记忆查询  
query = x + h_attn  # 破坏查询精准性，导致记忆选择固化
query = x  # 缺乏上下文处理，查询精度不足

记忆选择多样性监控：

• 定期检查不同输入位置的记忆选择分布
• 监控记忆选择熵值，确保 > 2.0
• 避免所有位置都选择相同记忆条目的情况

训练健康性检查：

• 训练Loss不应过快下降到极低值 (<0.1)
• 定期进行自回归推理评估，防止记忆化学习
• 训练-推理Loss差异应保持在合理范围内 (<10倍)

常见问题

1. 文本生成质量问题

• 现象: 输出为词组碎片，缺乏连贯性
• 可能原因: KnowledgeDataset 记忆机制与语言建模目标不匹配
• 排查方向: 检查知识库索引机制、记忆层输出分布

2. SFT 效果差异

• 现象: model 的 SFT 效果显著低于 baseline
• 可能原因: 预训练阶段的表示学习偏差
• 排查方向: 对比两种模型的隐层表示、梯度流动

3. 灾难性过拟合 (新增)

• 现象: 训练Loss极低但推理Loss极高，生成文本完全重复
• 根本原因: 查询机制破坏导致记忆选择固化
• 预防措施: 严格遵循查询特异性原则，实施早期预警监控

4. 训练资源

• GPU 内存: 如遇显存不足，调整 batch_size / accumulation_steps
• 训练速度: 确认 DeepSpeed ZeRO Stage 2 正确启用

调试工具

# 检查模型加载
.venv/bin/python -c "from model.model import *; print('模型加载成功')"

# 验证数据预处理
.venv/bin/python -c "from model.dataset import *; print('数据集加载成功')"

# 测试训练脚本
.venv/bin/python train_pretrain_accelerate.py --help

# 测试评估脚本
.venv/bin/python eval_model.py --help

# 快速评估测试（仅5个样本）
.venv/bin/python eval_model.py \
    --model_path out/experiment_1_4_0/pretrain_512.pth \
    --model_type model \
    --num_samples 5

---

💡 提示: 使用本文档前，请确保已正确配置 uv 虚拟环境和相关依赖。如有问题，请检查 pyproject.toml 配置。