From 65c89325fa57e6f35b7993456f91109f7944790c Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 11 Apr 2026 13:44:07 +0800 Subject: [PATCH] first commit --- .gitignore | 0 main.py | 305 ++++++++++++++++++++++++++++++++++++++ video.md | 426 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 731 insertions(+) create mode 100644 .gitignore create mode 100644 main.py create mode 100644 video.md diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e69de29 diff --git a/main.py b/main.py new file mode 100644 index 0000000..5390da9 --- /dev/null +++ b/main.py @@ -0,0 +1,305 @@ +""" +AstrAI项目介绍视频 - 第一部分：开场白与项目定位 +""" +from manim import * + + +class Part1_Intro(Scene): + """第一部分：开场白与项目定位""" + + def construct(self): + # === 场景1：Logo动画 === + self.play_logo_intro() + + # === 场景2：开场白 === + self.play_opening_speech() + + # === 场景3：核心要点展示 === + self.play_key_points() + + # === 场景4：项目名称与总结 === + self.play_project_name() + + def play_logo_intro(self): + """Logo动画展示""" + # 创建AstrAI Logo样式 + logo_text = Text( + "AstrAI", + font_size=72, + font="SimHei", + ) + + # 创建副标题 + subtitle = Text( + "轻量级 Transformer 训练推理框架", + font_size=28, + font="SimHei", + color=GRAY + ) + + # 居中排列 + logo_text.next_to(subtitle, UP, buff=0.3) + subtitle.next_to(logo_text, DOWN, buff=0.3) + + # 初始状态 + logo_text.set_opacity(0) + subtitle.set_opacity(0) + + # 动画：淡入Logo + self.play( + FadeIn(logo_text, shift=UP * 0.5, scale=1.2), + run_time=1.5 + ) + + # 动画：淡入副标题 + self.play( + FadeIn(subtitle, shift=UP * 0.3), + run_time=1.0 + ) + + # 短暂停留 + self.wait(0.8) + + # 淡出 + self.play( + FadeOut(logo_text), + FadeOut(subtitle), + run_time=0.8 + ) + self.wait(0.5) + + def play_opening_speech(self): + """开场白""" + # 第一段文字 + line1 = Text( + "在当下的大语言模型时代，", + font_size=32, + font="SimHei" + ) + line1.to_edge(UP, buff=1.0) + + self.play(Write(line1), run_time=1.0) + self.wait(0.5) + + # 第二段文字 + line2 = Text( + "我们见证了GPT、LLaMA等庞然大物。", + font_size=32, + font="SimHei" + ) + line2.next_to(line1, DOWN, buff=0.3) + + self.play(Write(line2), run_time=1.0) + self.wait(0.5) + + # 第三段文字 + line3 = Text( + "然而，这些模型对硬件要求极高，", + font_size=32, + font="SimHei" + ) + line3.next_to(line2, DOWN, buff=0.3) + + self.play(Write(line3), run_time=1.0) + self.wait(0.5) + + # 第四段文字（关键问题） + line4 = Text( + "对于普通开发者来说往往遥不可及。", + font_size=32, + font="SimHei", + color=BLUE + ) + line4.next_to(line3, DOWN, buff=0.3) + + self.play(Write(line4), run_time=1.0) + self.wait(0.8) + + # 思考问题 + question = Text( + "我们能否创建一个既有用又能运行在普通电脑上的模型？", + font_size=28, + font="SimHei", + color=YELLOW + ) + question.to_edge(DOWN, buff=1.5) + + self.play(Write(question), run_time=1.5) + self.wait(1.0) + + # 清除场景 + self.play( + FadeOut(line1), + FadeOut(line2), + FadeOut(line3), + FadeOut(line4), + FadeOut(question), + run_time=0.8 + ) + self.wait(0.5) + + def play_key_points(self): + """核心要点展示""" + # 标题 + title = Text( + "AstrAI 的核心特点", + font_size=40, + font="SimHei", + color=WHITE + ) + title.to_edge(UP, buff=0.8) + + self.play(Write(title), run_time=1.0) + self.wait(0.5) + + # 要点1：轻量级 + point1_title = Text( + "轻量级", + font_size=36, + font="SimHei", + color=GREEN + ) + point1_title.to_edge(LEFT, buff=2.0).shift(UP * 0.5) + + point1_content = Text( + "1B参数，中英双语", + font_size=24, + font="SimHei", + color=GRAY + ) + point1_content.next_to(point1_title, DOWN, buff=0.2) + point1_content.align_to(point1_title, LEFT) + + self.play( + Write(point1_title), + Write(point1_content), + run_time=1.0 + ) + self.wait(0.8) + + # 要点2：完全自研 + point2_title = Text( + "完全自研", + font_size=36, + font="SimHei", + color=BLUE + ) + point2_title.next_to(point1_title, DOWN, buff=0.8) + point2_title.align_to(point1_title, LEFT) + + point2_content = Text( + "核心代码全部开源", + font_size=24, + font="SimHei", + color=GRAY + ) + point2_content.next_to(point2_title, DOWN, buff=0.2) + point2_content.align_to(point2_title, LEFT) + + self.play( + Write(point2_title), + Write(point2_content), + run_time=1.0 + ) + self.wait(0.8) + + # 要点3：训练推理一体化 + point3_title = Text( + "训练推理一体化", + font_size=36, + font="SimHei", + color=ORANGE + ) + point3_title.next_to(point2_title, DOWN, buff=0.8) + point3_title.align_to(point1_title, LEFT) + + point3_content = Text( + "从预训练到推理，一套全搞定", + font_size=24, + font="SimHei", + color=GRAY + ) + point3_content.next_to(point3_title, DOWN, buff=0.2) + point3_content.align_to(point3_title, LEFT) + + self.play( + Write(point3_title), + Write(point3_content), + run_time=1.0 + ) + self.wait(1.0) + + # 清除场景（保留标题） + self.play( + FadeOut(point1_title), + FadeOut(point1_content), + FadeOut(point2_title), + FadeOut(point2_content), + FadeOut(point3_title), + FadeOut(point3_content), + run_time=0.5 + ) + self.wait(0.3) + + def play_project_name(self): + """项目名称与总结""" + # 项目名称 - 使用普通颜色避免兼容性问题 + project_name = Text( + "AstrAI", + font_size=80, + font="SimHei", + color=BLUE + ) + project_name.center() + + # 动画：放大出现 + self.play( + project_name.animate.scale(1.2), + run_time=1.5 + ) + self.wait(1.0) + + # 补充说明 + tagline = Text( + "1B参数 · 中英双语 · 完全自研 · 训练推理一体化", + font_size=24, + font="SimHei", + color=YELLOW + ) + tagline.next_to(project_name, DOWN, buff=0.5) + + self.play(Write(tagline), run_time=1.0) + self.wait(1.0) + + # 淡出结束 + self.play( + FadeOut(project_name), + FadeOut(tagline), + run_time=1.0 + ) + self.wait(0.5) + + +# 配置类 +class AstrAIVideoConfig: + """视频配置""" + video_quality = "high_quality" # high_quality, production_quality + preview = True # 是否预览 + save_last_frame = True # 保存最后一帧 + + +if __name__ == "__main__": + # 运行配置 + config = AstrAIVideoConfig() + + # 渲染视频 + if config.preview: + # 预览模式 + from manim import * + scene = Part1_Intro() + scene.render() + else: + # 渲染视频文件 + # manim -pql main.py Part1_Intro + # manim -pqh main.py Part1_Intro + pass \ No newline at end of file diff --git a/video.md b/video.md new file mode 100644 index 0000000..f58600e --- /dev/null +++ b/video.md @@ -0,0 +1,426 @@ +# 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 仓���** +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 \ No newline at end of file