LLM-lite: A Lightweight Framework for LLM Exploratio

项目概述

完整的 DeepSeek-V3 等先进大语言模型涉及到庞大的工程优化和算力资源。本项目旨在提供一个轻量级的框架，专注于：

1. 核心模型架构学习: 在有限算力下，复现主流 LLM（如 DeepSeek-V3, Llama3）的核心机制，如 Attention, MoE, RoPE 等。
2. 可复用的训练与推理: 构建一个易于扩展的训练和推理流程，方便快速迭代和学习新模型。
3. 教学与实验: 为希望深入理解 LLM 底层原理的学习者和研究人员提供一个清晰、简洁的代码基础。

通过本框架，用户可以将主要精力聚焦在模型架构本身，而非复杂的工程细节。

主要特性

• 模块化模型设计: 基于 BaseModel 类，方便添加和切换不同的 LLM 架构。
• 核心组件实现: 包含了 Attention (MLA for DeepSeekV3), MoE, MLP, RMSNorm, RoPE 等关键组件的 PyTorch 实现。
• 基础训练流程: 支持使用 PyTorch DDP进行分布式预训练，包含学习率调度、数据加载等。
• 基础推理功能: 支持生成文本，可配置 temperature, top-k, top-p 等参数。
• 分词器支持: 包含预训练分词器，并提供训练脚本。
• 简洁易懂: 代码力求清晰，注释适中，便于学习和理解。

项目结构

LLM-lite/
├── configs/                # 配置文件 (模型参数, 训练设置等)
│   ├── deepseekv3_small.yaml
│   └── llama3_small.yaml
├── data/                   # 存放原始或处理后的数据集 (建议添加到 .gitignore)
│   └── your_dataset/
├── outputs/                # 存放训练输出 (模型 checkpoints, 日志等 - 建议添加到 .gitignore)
│   └── runs/
├── tokenizer_data/         # 预训练的分词器文件
│   ├── merges.txt
│   ├── tokenizer.json
│   ├── tokenizer_config.json
│   └── vocab.json
├── llm_lite/               # 核心源代码包
│   ├── __init__.py
│   ├── models/             # 模型定义
│   │   ├── __init__.py
│   │   ├── base.py         # 模型基类
│   │   ├── deepseekv3.py
│   │   └── llama3.py
│   ├── data/               # 数据加载与处理代码
│   │   ├── __init__.py
│   │   └── datasets.py
│   ├── tokenizer/          # 分词器相关代码
│   │   ├── __init__.py
│   │   └── train.py
│   ├── utils/              # 工具函数
│   │   ├── __init__.py
│   │   └── tools.py
│   └── trainers/           # 训练逻辑 (可选)
│       ├── __init__.py
│       └── pretrainer.py
├── scripts/                # 可执行脚本
│   ├── pretrain.py         # 启动预训练
│   ├── generate.py         # 文本生成/推理
│   └── evaluate.py         # (未来) 模型评估脚本
├── requirements.txt
├── README.md
└── .gitignore

快速开始

1. 克隆仓库:

   git clone https://liuruing/LLM-lite.git
   cd LLM-lite

2. 安装依赖:

   pip install -r requirements.txt

3. 准备数据: 将训练数据放入 data/ 目录。
4. 配置模型和训练: 在 configs/ 目录下创建或修改 YAML 配置文件。
5. 训练模型: (脚本路径已更改)

   # 示例: 使用配置文件启动训练
   python scripts/pretrain.py --config_path configs/llama3_small.yaml
   
   # 或者仍然支持命令行参数覆盖
   torchrun --standalone --nproc_per_node=2 scripts/pretrain.py --config_path configs/deepseekv3_small.yaml --batch_size 8

6. 运行推理: (脚本路径和名称已更改)

   python scripts/generate.py --config_path configs/llama3_small.yaml --checkpoint_path outputs/runs/.../checkpoint.pt
   # 或者如果是 Streamlit 应用 (假设脚本已更新)
   # streamlit run scripts/generate.py ...

使用方法

模型训练 (scripts/pretrain.py)

• 配置文件: 推荐使用 configs/ 下的 YAML 文件来管理所有模型和训练参数。通过 --config_path 指定。
• 命令行覆盖: 命令行参数可以覆盖配置文件中的设置。
• 核心逻辑: (可选) 训练的核心循环可能在 llm_lite/trainers/pretrainer.py 中实现。
• 分布式训练: 同样使用 torchrun 配合 scripts/pretrain.py 启动。

模型推理 (scripts/generate.py)

• 使用 --config_path 加载模型结构配置。
• 使用 --checkpoint_path 指定训练好的模型权重。
• 调用 llm_lite.models.base.BaseModel.generate 方法进行生成。
• 可配置 max_new_tokens, temperature, top_k, top_p 等推理参数。

添加新模型

1. 创建模型文件: 在 llm_lite/models/ 目录下创建新的 Python 文件 (例如 mymodel.py)。
2. 定义模型参数: 在 mymodel.py 中定义继承自 llm_lite.models.base.BaseModelArgs 的 dataclass。
3. 实现模型架构: 定义继承自 llm_lite.models.base.BaseModel 的新模型类，实现 __init__ 和 forward，设置 model_name。
4. 注册模型: 在 llm_lite/models/__init__.py 中导入并添加到 MODELS 字典。
5. 创建配置文件: 在 configs/ 目录下为新模型创建一个 YAML 配置文件。
6. 现在可以在 scripts/pretrain.py 或 scripts/generate.py 中通过 --config_path 使用你的新模型了。

已实现模型

• mini_llama3: Llama3 的简化版本。
• mini_deepseekv3: DeepSeek-V3 的简化版本，包含 MLA 和 MoE 的基础实现。

贡献与未来工作

欢迎贡献代码、报告问题或提出改进建议！

• TODO: 添加更详细的 MTP 实现说明。
• TODO: 增加模型评估脚本。
• TODO: 探索更高效的训练策略。
• TODO: 支持加载预训练权重。

---