426 lines
12 KiB
Markdown
426 lines
12 KiB
Markdown
# AstrAI 项目介绍视频文稿
|
||
|
||
## 📖 目录
|
||
|
||
- [一、核心论点与总览](#一核心论点与总览)
|
||
- [1.1 开场白与项目定位](#11-开场白与项目定位)
|
||
- [二、AstrAI的核心设计](#二astrai的核心设计)
|
||
- [2.1 整体架构](#21-整体架构)
|
||
- [2.2 推理流程：Prefill → Decode](#22-推理流程prefill--decode)
|
||
- [2.3 调度策略：高效管理KV Cache](#23-调度策略高效管理kv-cache)
|
||
- [2.4 内存优化：对比vLLM的PagedAttention](#24-内存优化对比vllm的pagedattention)
|
||
- [2.5 分布式支持：多卡推理](#25-分布式支持多卡推理)
|
||
- [三、为什么选择自研？](#三为什么选择自研)
|
||
- [3.1 主流框架对比](#31-主流框架对比)
|
||
- [3.2 各框架详细对比](#32-各框架详细对比)
|
||
- [四、AstrAI的局限性](#四astrai的局限性)
|
||
- [4.1 功能生态](#41-功能生态)
|
||
- [4.2 社区和文档](#42-社区和文档)
|
||
- [4.3 性能验证](#43-性能验证)
|
||
- [五、未来规划与贡献方式](#五未来规划与贡献方式)
|
||
- [5.1 未来规划](#51-未来规划)
|
||
- [5.2 如何参与](#52-如何参与)
|
||
- [六、快速开始](#六快速开始)
|
||
- [6.1 安装](#61-安装)
|
||
- [6.2 训练模型](#62-训练模型)
|
||
- [6.3 文本生成](#63-文本生成)
|
||
- [6.4 启动推理服务](#64-启动推理服务)
|
||
- [七、总结](#七总结)
|
||
- [附录：核心代码示例](#附录核心代码示例)
|
||
|
||
---
|
||
|
||
## 视频脚本
|
||
|
||
---
|
||
|
||
### 一、核心论点与总览
|
||
|
||
#### 1.1 开场白与项目定位
|
||
|
||
> 【画面建议：项目Logo动画，展示AstrAI的品牌形象】
|
||
|
||
大家好，今天要给大家介绍的是一个完全自研的轻量级Transformer训练推理一体化框架——AstrAI。
|
||
|
||
在当下的大语言模型时代，我们见证了GPT、LLaMA等拥有数十亿甚至数百亿参数的庞然大物。然而，这些模型对硬件的要求极高，对于普通开发者来说往往遥不可及。我一直在思考：**我们能否创建一个既有用又能运行在普通电脑上的模型？**
|
||
|
||
这也正是大多数人所期望的——一个本地可部署的AI项目，在实现完全私有化的同时，保持一定的智能水平。
|
||
|
||
> 【要点强调】
|
||
> - **轻量级**：1B参数，中英双语
|
||
> - **完全自研**：核心代码全部开源
|
||
> - **训练推理一体化**：从预训练到推理，一套框架全搞定
|
||
|
||
于是，AstrAI项目应运而生——1B参数、中英双语、支持对话和文本生成，训练代码完全开源！
|
||
|
||
---
|
||
|
||
### 二、AstrAI的核心设计
|
||
|
||
#### 2.1 整体架构
|
||
|
||
> 【画面建议：架构图动画，展示各个模块之间的关系】
|
||
|
||
AstrAI的系统架构采用了模块化设计，主要包含以下核心模块：
|
||
|
||
| 模块 | 描述 |
|
||
|------|------|
|
||
| **astrai.config** | ModelConfig, TrainConfig - 配置管理 |
|
||
| **astrai.dataset** | DatasetFactory - 数据集加载与管理，支持SEQ/SFT/DPO/GRPO |
|
||
| **astrai.model** | AutoModel, Transformer - 神经网络模型 |
|
||
| **astrai.tokenize** | AutoTokenizer, ChatTemplate - 分词器和聊天模板 |
|
||
| **astrai.trainer** | Trainer - 训练工作流管理 |
|
||
| **astrai.inference** | InferenceScheduler - 推理调度，支持连续批处理 |
|
||
| **astrai.parallel** | 分布式并行支持 |
|
||
|
||
---
|
||
|
||
#### 2.2 推理流程：Prefill → Decode
|
||
|
||
> 【画面建议：Pre-fill和Decode阶段的流程动画】
|
||
|
||
推理是AstrAI的核心能力之一。整个推理过程分为两个阶段：
|
||
|
||
**第一阶段：Prefill（预填充）**
|
||
|
||
在Prefill阶段，系统一次性处理整个输入序列，计算出所有token的Key和Value，并向量化。这个阶段是计算密集型的，需要处理完整的输入。
|
||
|
||
**第二阶段：Decode（解码）**
|
||
|
||
在Decode阶段，系统只需要为新生成的token计算Query，并从缓存的KV Cache中获取之前所有token的Key和Value进行注意力计算。这就是KV Cache的威力所在——大幅减少重复计算。
|
||
|
||
> 【技术细节】
|
||
> 根据注意力计算公式：
|
||
> $$ o_n = \sum_j \text{softmax}\left(\frac{q_n k_{j}}{\sqrt{d_k}}\right)v_{j} $$
|
||
> 其中只有k和v有长度索引，q固定为最后一个token。因此我们需要缓存k和v以避免重复计算。
|
||
> 需要注意的是：**位置编码必须在KV Cache计算之前进行**，否则会出现位置编码错误。
|
||
|
||
---
|
||
|
||
#### 2.3 调度策略：高效管理KV Cache
|
||
|
||
> 【画面建议：任务调度队列和管理器的可视化】
|
||
|
||
AstrAI使用InferenceScheduler来管理多个推理任务，它采用生产者-消费者模式：
|
||
|
||
```python
|
||
class InferenceScheduler:
|
||
def __init__(
|
||
self,
|
||
model: AutoModel,
|
||
tokenizer: AutoTokenizer,
|
||
max_batch_size: int = 16,
|
||
max_seq_len: Optional[int] = None,
|
||
max_prefix_len: int = 512,
|
||
cache_capacity: int = 1000,
|
||
):
|
||
self.waiting_queue: List[Task] = [] # 等待中的任务队列
|
||
self.active_tasks: List[Task] = [] # 正在执行的任务列表
|
||
self.kv_cache = (k_cache, v_cache) # KV Cache
|
||
self.prefix_cache = PrefixCacheManager(max_capacity=cache_capacity) # 前缀缓存
|
||
...
|
||
|
||
def add_task(): 添加新任务
|
||
def remove_task(): 移除完成任务
|
||
def start(): 启动推理循环
|
||
def stop(): 停止推理并释放资源
|
||
```
|
||
|
||
**连续批处理**是AstrAI的特色功能：
|
||
- 动态批处理：新的请求可以随时加入
|
||
- 立即释放：完成的任务立即释放资源
|
||
- 这大大提高了GPU利用率
|
||
- **前缀缓存**：使用Radix Tree实现智能前缀提示（Prefix Caching）加速
|
||
|
||
---
|
||
|
||
#### 2.4 内存优化：对比vLLM的PagedAttention
|
||
|
||
> 【画面建议：KV Cache内存管理机制的对比图】
|
||
|
||
为了优化内存使用，AstrAI实现了PrefixCacheManager，使用**Radix Tree（前缀树）**来管理KV Cache：
|
||
|
||
```python
|
||
class RadixNode:
|
||
- children: Dict[int, RadixNode] # token_id -> 子节点
|
||
- hash: Optional[int] # 前缀的64位哈希
|
||
- slot: int # KV Cache槽位
|
||
- ref_count: int # 引用计数
|
||
- last_access: float # 最近访问时间戳
|
||
|
||
class PrefixCacheManager:
|
||
- root: RadixNode # Radix树根节点
|
||
- max_capacity: 最大缓存容量
|
||
- lru: List[Tuple[float, RadixNode]] # LRU队列
|
||
- insert(): 插入新token序列
|
||
- find_longest_prefix(): 查找最长前缀匹配
|
||
- release(): 释放缓存
|
||
```
|
||
|
||
**Radix Tree的优势**：
|
||
- 自动合并共享前缀的请求
|
||
- 智能复用已计算的KV Cache
|
||
- 支持前缀提示（Prefix Caching）加速
|
||
- 基于LRU的缓存淘汰策略
|
||
|
||
> 【对比vLLM】
|
||
> vLLM使用PagedAttention进行分页内存管理，而AstrAI使用Radix Tree实现前缀缓存。两者都能有效减少内存碎片，但实现方式不同。AstrAI的优势在于轻量级和训练推理一体化。
|
||
|
||
---
|
||
|
||
#### 2.5 分布式支持：多卡推理
|
||
|
||
> 【画面建议：多卡并行的架构示意图】
|
||
|
||
AstrAI支持多卡分布式推理，使用以下并行策略：
|
||
|
||
```python
|
||
from astrai.parallel import setup_parallel, spawn_parallel_fn
|
||
|
||
# 初始化分布式环境
|
||
with setup_parallel(
|
||
rank=0,
|
||
world_size=4,
|
||
backend="nccl",
|
||
device_type="cuda",
|
||
):
|
||
# 在分布式环境中执行推理
|
||
...
|
||
|
||
# 启动多进程分布式推理
|
||
spawn_parallel_fn(
|
||
func=inference_fn,
|
||
world_size=4,
|
||
backend="nccl",
|
||
device_type="cuda",
|
||
)
|
||
```
|
||
|
||
支持的并行方式：
|
||
- **数据并行**：多卡处理不同数据
|
||
- **模型并行**：将模型分片到不同卡（使用RowParallelLinear, ColumnParallelLinear）
|
||
- **流水并行**：按层划分Pipeline
|
||
|
||
---
|
||
|
||
### 三、为什么选择自研？
|
||
|
||
#### 3.1 主流框架对比
|
||
|
||
> 【画面建议：三种框架的对比表】
|
||
|
||
| 特性 | HuggingFace Transformers | vLLM | **AstrAI** |
|
||
|------|-------------------------|------|------------|
|
||
| **定位** | 通用深度学习框架 | 专用推理框架 | 训练推理一体化 |
|
||
| **代码量** | 大而全 | 中等 | 轻量级 |
|
||
| **依赖** | 众多 | 中等 | 最小化 |
|
||
| **训练支持** | 仅模型训练 | 不支持 | 完整训练流程 |
|
||
| **推理优化** | 基础 | 高级优化 | 中等优化 |
|
||
| **连续批处理** | 不支持 | 支持 | 支持 |
|
||
|
||
---
|
||
|
||
#### 3.2 各框架详细对比
|
||
|
||
**HuggingFace Transformers：灵活但重**
|
||
|
||
HuggingFace Transformers是目前最流行的深度学习框架，提供了海量的预训练模型和工具。它的优势在于：
|
||
- 生态丰富，模型众多
|
||
- API设计优雅，易于使用
|
||
- 社区活跃，文档完善
|
||
|
||
但它的问题也很明显：
|
||
- 代码库庞大，依赖众多
|
||
- 对推理的优化不够深入
|
||
- 不支持连续批处理
|
||
|
||
**vLLM：专注推理但功能单一**
|
||
|
||
vLLM是专门为推理优化的框架，PagedAttention技术独领风骚：
|
||
- 内存效率高，吞吐量出众
|
||
- 支持连续批处理
|
||
- 推理性能业界领先
|
||
|
||
但它只专注推理，不支持训练。
|
||
|
||
**AstrAI：轻量 + 训练推理一体化**
|
||
|
||
AstrAI的定位清晰：
|
||
- **轻量级**：最小化依赖，易于部署
|
||
- **训练推理一体化**：从预训练到推理全覆盖
|
||
- **研究友好**：模块化设计，易于实验新想法
|
||
- 中国开发者友好：中文文档，本地支持
|
||
|
||
---
|
||
|
||
### 四、AstrAI的局限性
|
||
|
||
> 【画面建议：坦诚面对局限性的态度】
|
||
|
||
任何项目都有其局限性，AstrAI也不例外。
|
||
|
||
#### 4.1 功能生态
|
||
|
||
- 功能生态不如vLLM完善
|
||
- 不支持一些高级推理特性
|
||
- 插件和扩展较少
|
||
|
||
#### 4.2 社区和文档
|
||
|
||
- 社区和文档还在建设中
|
||
- 用户基数相对较小
|
||
- 响应速度可能不如主流框架
|
||
|
||
#### 4.3 性能验证
|
||
|
||
- 性能还需要更多benchmark验证
|
||
- 缺少大规模生产环境验证
|
||
- 与顶级框架的性能对比数据有限
|
||
|
||
> 【态度】
|
||
> 这些都是我们正在努力改进的方向。我们真诚地欢迎社区的参与和贡献！
|
||
|
||
---
|
||
|
||
### 五、未来规划与贡献方式
|
||
|
||
#### 5.1 未来规划
|
||
|
||
> 【画面建议：路线图展示】
|
||
|
||
- 完善文档和教程
|
||
- 增加更多模型架构支持
|
||
- 性能优化和benchmark
|
||
- 丰富的预训练数据集
|
||
- 更强的社区支持
|
||
|
||
#### 5.2 如何参与
|
||
|
||
> 【画面建议：GitHub页面展示】
|
||
|
||
**欢迎Star和Fork！**
|
||
|
||
```bash
|
||
# 克隆项目
|
||
git clone https://github.com/ViperEkura/AstrAI.git
|
||
cd AstrAI
|
||
pip install -e .
|
||
```
|
||
|
||
**贡献方式：**
|
||
|
||
1. **Fork 仓<><E4BB93><EFBFBD>**
|
||
2. **创建特性分支**
|
||
3. **提交你的修改**
|
||
4. **发起Pull Request**
|
||
|
||
对于重大修改，请先提Issue讨论再动手。
|
||
|
||
---
|
||
|
||
### 六、快速开始
|
||
|
||
#### 6.1 安装
|
||
|
||
```bash
|
||
git clone https://github.com/ViperEkura/AstrAI.git
|
||
cd AstrAI
|
||
pip install -e .
|
||
```
|
||
|
||
#### 6.2 训练模型
|
||
|
||
```bash
|
||
python scripts/tools/train.py \
|
||
--train_type=seq \
|
||
--data_root_path=/path/to/dataset \
|
||
--param_path=/path/to/param_path
|
||
```
|
||
|
||
#### 6.3 文本生成
|
||
|
||
```bash
|
||
python scripts/tools/generate.py --param_path=/path/to/param_path
|
||
```
|
||
|
||
#### 6.4 启动推理服务
|
||
|
||
```bash
|
||
python -m scripts.tools.server --port 8000 --device cuda
|
||
```
|
||
|
||
---
|
||
|
||
### 七、总结
|
||
|
||
> 【画面建议：项目Logo和口号】
|
||
|
||
AstrAI，一个**轻量级**、**完全自研**、**训练推理一体化**的Transformer框架。
|
||
|
||
它让大模型不再高高在上，让每个开发者都能在普通设备上运行自己的AI。
|
||
|
||
**行动号召：**
|
||
- 给项目加个Star：https://github.com/ViperEkura/AstrAI
|
||
- 欢迎提Issue和参与讨论
|
||
- 一起打造中国开发者的AI框架！
|
||
|
||
---
|
||
|
||
## 附录：核心代码示例
|
||
|
||
### AutoModel使用
|
||
|
||
```python
|
||
from astrai.model import AutoModel
|
||
|
||
# 加载模型
|
||
model = AutoModel.from_pretrained("path/to/model")
|
||
|
||
# 保存模型
|
||
model.save_pretrained("path/to/save")
|
||
```
|
||
|
||
### 推理引擎使用
|
||
|
||
```python
|
||
from astrai.inference import InferenceEngine, GenerationRequest
|
||
from astrai.model import AutoModel
|
||
from astrai.tokenize import AutoTokenizer
|
||
|
||
# 加载模型和分词器
|
||
model = AutoModel.from_pretrained("path/to/model")
|
||
tokenizer = AutoTokenizer.from_pretrained("path/to/model")
|
||
|
||
# 创建推理引擎
|
||
engine = InferenceEngine(
|
||
model=model,
|
||
tokenizer=tokenizer,
|
||
max_batch_size=8,
|
||
max_seq_len=4096,
|
||
)
|
||
|
||
# 创建请求
|
||
request = GenerationRequest(
|
||
messages=[
|
||
{"role": "system", "content": "You are a helpful assistant."},
|
||
{"role": "user", "content": "Hello"},
|
||
],
|
||
temperature=0.8,
|
||
top_p=0.95,
|
||
top_k=50,
|
||
max_len=1024,
|
||
stream=True,
|
||
)
|
||
|
||
# 流式生成
|
||
for token in engine.generate_with_request(request):
|
||
print(token, end="", flush=True)
|
||
|
||
# 关闭引擎
|
||
engine.shutdown()
|
||
```
|
||
|
||
---
|
||
|
||
> **视频文稿版本**：1.0
|
||
> **更新时间**：2026-04-11 |