253 lines
7.2 KiB
Markdown
253 lines
7.2 KiB
Markdown
# MiniMind 预训练项目开发文档
|
||
|
||
## 项目概述
|
||
|
||
MiniMind 是一个基于 Transformer 架构的大语言模型预训练项目，集成了先进的知识图谱技术和混合专家模型（MOE）架构。项目采用 PyTorch 实现，支持分布式训练和高效的内存管理。
|
||
|
||
## 核心架构
|
||
|
||
### 1. 主训练入口
|
||
|
||
**`train_pretrain_accelerate.py`** - 主训练脚本，包含完整的训练流程：
|
||
|
||
- **内存监控系统**: 实时监控系统内存和 GPU 内存使用情况
|
||
- **分布式训练**: 基于 Accelerate 和 DeepSpeed 的分布式训练支持
|
||
- **知识库初始化**: 从 JSON 数据文件初始化知识库，支持缓存机制
|
||
- **训练循环**: 包含梯度累积、学习率调度、损失计算等完整训练逻辑
|
||
|
||
### 2. 模型架构
|
||
|
||
**`model/model.py`** - 核心模型实现：
|
||
|
||
```python
|
||
class MiniMindLM(PreTrainedModel):
|
||
"""主要的 Transformer 模型类"""
|
||
- 标准 Transformer 架构（decoder-only）
|
||
- RMSNorm 归一化层
|
||
- 旋转位置编码（RoPE）
|
||
- Flash Attention 支持
|
||
- 知识库集成
|
||
```
|
||
|
||
**`model/LMConfig.py`** - 模型配置类：
|
||
|
||
```python
|
||
class LMConfig(PretrainedConfig):
|
||
"""模型配置管理"""
|
||
- 基础模型参数（dim, n_layers, n_heads 等）
|
||
- MOE 相关配置
|
||
- 知识图谱配置
|
||
- 数据库功能配置
|
||
```
|
||
|
||
### 3. 知识库系统
|
||
|
||
**`KnowledgeDataset`** 类（在 `model/model.py` 中）：
|
||
|
||
- **二维分解键空间**: 使用 Product Key 方法优化大规模知识库检索
|
||
- **智能选择策略**: 动态调整知识库访问模式
|
||
- **可训练参数**: 键向量支持梯度更新
|
||
- **缓存机制**: 支持知识库预处理结果缓存
|
||
|
||
### 4. 数据处理
|
||
|
||
**`model/dataset.py`** - 数据集处理：
|
||
|
||
```python
|
||
class PretrainDataset(Dataset):
|
||
"""预训练数据集类"""
|
||
- JSONL 格式数据加载
|
||
- 自动添加 BOS/EOS 标记
|
||
- 序列填充和截断
|
||
- 损失掩码生成
|
||
```
|
||
|
||
## 核心功能模块
|
||
|
||
### 1. 内存管理
|
||
|
||
项目实现了完善的内存监控系统：
|
||
|
||
```python
|
||
def get_memory_usage():
|
||
"""获取系统内存使用情况"""
|
||
|
||
def get_cuda_memory_usage():
|
||
"""获取 GPU 内存使用情况"""
|
||
|
||
def log_memory_status():
|
||
"""记录详细的内存状态"""
|
||
```
|
||
|
||
### 2. 知识库初始化
|
||
|
||
知识库初始化流程：
|
||
|
||
1. **数据加载**: 从 JSON 文件加载句子数据
|
||
2. **重要性排序**: 根据 importance_score 对句子排序
|
||
3. **分词处理**: 使用 tokenizer 将句子转换为 token 序列
|
||
4. **长度处理**: 截断或填充到指定长度
|
||
5. **缓存机制**: 支持处理结果缓存以加速后续训练
|
||
|
||
### 3. 分布式训练配置
|
||
|
||
**Accelerate 配置** (`accelerate_config.yaml`):
|
||
```yaml
|
||
compute_environment: LOCAL_MACHINE
|
||
distributed_type: DEEPSPEED
|
||
mixed_precision: bf16
|
||
num_processes: 4
|
||
deepspeed_config:
|
||
deepspeed_config_file: ds_config.json
|
||
```
|
||
|
||
**DeepSpeed 配置** (`ds_config.json`):
|
||
```json
|
||
{
|
||
"zero_optimization": {
|
||
"stage": 2,
|
||
"offload_optimizer": {
|
||
"device": "cpu",
|
||
"pin_memory": true
|
||
}
|
||
},
|
||
"optimizer": {
|
||
"type": "AdamW"
|
||
},
|
||
"scheduler": {
|
||
"type": "WarmupLR"
|
||
}
|
||
}
|
||
```
|
||
|
||
## 主要配置参数
|
||
|
||
### 模型配置
|
||
- `dim`: 隐藏层维度（默认 512）
|
||
- `n_layers`: Transformer 层数（默认 8）
|
||
- `n_heads`: 注意力头数（默认 32）
|
||
- `n_kv_heads`: KV 注意力头数（默认 8）
|
||
- `max_seq_len`: 最大序列长度（默认 512）
|
||
- `vocab_size`: 词汇表大小（默认 6400）
|
||
|
||
### 知识库配置
|
||
- `knowledge_num`: 知识库条目数量（默认 1048576）
|
||
- `knowledge_length`: 每个知识条目的长度（默认 32）
|
||
- `knowledge_dim`: 知识向量维度（默认 128）
|
||
|
||
### 训练配置
|
||
- `batch_size`: 批次大小（默认 128）
|
||
- `learning_rate`: 学习率（默认 8e-5）
|
||
- `accumulation_steps`: 梯度累积步数（默认 16）
|
||
- `warmup_iters`: 预热迭代次数
|
||
|
||
## 数据格式
|
||
|
||
### 预训练数据格式
|
||
```json
|
||
{"text": "这是一个训练样本的文本内容"}
|
||
```
|
||
|
||
### 知识库数据格式
|
||
```json
|
||
[
|
||
{
|
||
"target": [
|
||
{
|
||
"sentence": "知识库中的句子内容",
|
||
"importance_score": 0.95
|
||
}
|
||
]
|
||
}
|
||
]
|
||
```
|
||
|
||
## 工具脚本
|
||
|
||
### 数据预处理脚本
|
||
- `preprocessing/preprocess_pretrain.py`: 预训练数据预处理
|
||
- `preprocessing/preprocess_trex.py`: 三元组数据预处理
|
||
- `preprocessing/preprocess_combined_json.py`: 组合数据预处理
|
||
|
||
### 模型工具
|
||
- `dataset_decoder.py`: 解码模型中的知识库内容
|
||
|
||
### 运行脚本
|
||
- `run_file/experiment_*.sh`: 各种实验配置的运行脚本
|
||
|
||
## 依赖管理
|
||
|
||
项目使用 `pyproject.toml` 管理依赖：
|
||
|
||
### 核心依赖
|
||
- `torch >= 2.7.1`: 深度学习框架
|
||
- `transformers >= 4.52.4`: Transformer 模型库
|
||
- `accelerate >= 1.7.0`: 分布式训练
|
||
- `deepspeed >= 0.17.0`: 深度学习优化
|
||
- `swanlab >= 0.6.4`: 实验监控
|
||
|
||
### 开发工具
|
||
- `tokenizers >= 0.21.1`: 高效分词
|
||
- `datasets >= 2.21.0`: 数据集处理
|
||
- `numpy >= 1.26.4`: 数值计算
|
||
- `pandas >= 2.0.0`: 数据处理
|
||
|
||
## 内存优化策略
|
||
|
||
1. **梯度累积**: 通过累积梯度减少内存占用
|
||
2. **混合精度训练**: 使用 bf16 减少内存使用
|
||
3. **ZeRO 优化**: DeepSpeed ZeRO Stage 2 优化器状态分片
|
||
4. **知识库缓存**: 预处理结果缓存避免重复计算
|
||
5. **垃圾回收**: 定期清理未使用的内存
|
||
|
||
## 监控和日志
|
||
|
||
### SwanLab 集成
|
||
- 训练损失监控
|
||
- 学习率变化追踪
|
||
- 内存使用情况记录
|
||
- 训练速度统计
|
||
|
||
### 日志系统
|
||
- 时间戳格式化输出
|
||
- 多进程日志同步
|
||
- 内存状态详细记录
|
||
- 训练进度追踪
|
||
|
||
## 目录结构详解
|
||
|
||
```
|
||
.
|
||
├── train_pretrain_accelerate.py # 主训练脚本
|
||
├── dataset_decoder.py # 知识库解码工具
|
||
├── model/ # 模型定义目录
|
||
│ ├── LMConfig.py # 模型配置类
|
||
│ ├── model.py # 主模型实现
|
||
│ ├── dataset.py # 数据集处理
|
||
│ ├── model_no_feed.py # 无反馈模型变体
|
||
│ ├── model_original.py # 原始模型变体
|
||
│ └── minimind_tokenizer/ # 分词器文件
|
||
├── preprocessing/ # 数据预处理脚本
|
||
├── run_file/ # 实验运行脚本
|
||
├── models/ # 模型检查点存储
|
||
├── accelerate_config.yaml # Accelerate 配置
|
||
├── ds_config.json # DeepSpeed 配置
|
||
├── pyproject.toml # 项目依赖配置
|
||
└── uv.lock # 依赖锁定文件
|
||
```
|
||
|
||
## 开发注意事项
|
||
|
||
1. **模型变体**: 项目包含多个模型变体，选择合适的模型类型
|
||
2. **知识库大小**: 根据可用内存调整知识库参数
|
||
3. **分布式配置**: 根据硬件配置调整并行参数
|
||
4. **缓存管理**: 合理使用缓存机制避免重复计算
|
||
5. **内存监控**: 关注内存使用情况，及时调整批次大小
|
||
|
||
## 扩展点
|
||
|
||
1. **新模型架构**: 通过继承 `PreTrainedModel` 实现新的模型变体
|
||
2. **数据处理**: 扩展 `PretrainDataset` 支持新的数据格式
|
||
3. **知识库优化**: 改进 `KnowledgeDataset` 的检索策略
|
||
4. **训练策略**: 在主训练循环中添加新的训练技巧
|
||
5. **监控扩展**: 集成更多监控指标和可视化工具 |