swiftLLM

斯威夫特法学硕士

该项目仍在开发中。某些功能可能尚未实现，并且文档可能不完整。

一个专为研究目的而定制的小型但功能强大的法学硕士推理系统。

仅 2k 行代码（vLLM 的 2%）即可获得与 vLLM 相当的性能。

为什么选择 SwiftLLM

用于 LLM 服务的开源框架有很多，包括 HuggingFace Transformers、vLLM、LightLLM、DistServe 和 DeepSpeed-MII。为什么选择 SwiftLLM？

原因是，这些框架是为生产量身定制的，而不是为研究量身定制的。它们配备了众多功能，例如 100 多个模型支持、各种硬件支持、LoRA、量化、多模态、前缀缓存、波束搜索等。虽然它是一种一体化的生产解决方案，但它们的代码库太大、太复杂，难以理解和修改（例如，vLLM 有 10 万多行代码），因此很难将它们用于研究目的。此外，他们的历史负担也是一个问题。

SwiftLLM 被设计为一个小型但功能强大的 LLM 推理系统，专为研究目的而定制。 “微小”意味着它只保留了研究所必需的功能，“强大”意味着它在性能上没有妥协，最后“迅速”意味着它易于理解和修改。在支持基本功能（参见下面的列表）并能够实现与 vLLM 相当的性能的同时，SwiftLLM 的代码库少于2k行代码（约 vLLM 的 2%），用 Python 和 OpenAI Triton（一种用于编写CUDA 内核），使其易于阅读、修改、调试、测试、扩展，并且可以轻松地与您新颖而精彩的研究想法集成。

功能列表

目前，SwiftLLM支持以下功能：

• 迭代调度和选择性批处理（Orca 中提出）
• PagedAttenton（vLLM 论文中提出）
• LLaMA / LLaMA2 / LLaMA3 模型（链接）及其变体
• 搭载预填充和解码（SARATHI 中提出）
• Flash Attention（在FlashAttention和FlashAttention-2中提出）
• 分页注意力 v2（也称为 Flash-Decoding，此处提出）

我们计划在未来添加对以下功能的支持：

• 张量并行和管道并行

为了保持代码库较小，我们将不支持以下功能。如果你想在你的研究项目中使用它们，你可能需要自己实现它们：

• 量化
• 洛拉
• 多式联运
• 不遵循 LLaMA 架构的模型
• 贪婪采样以外的采样方法
• 硬件支持除 NVIDIA GPU 之外的其他硬件（但只要 OpenAI Triton 支持，就应该很容易迁移到其他硬件）

请记住，SwiftLLM 并不是一款一体化的生产解决方案。建议将其视为您研究项目的“基础”，您可能需要自己实现一些功能。我亲爱的研究人员，我们鼓励您阅读代码、理解它、修改它并扩展它以满足您的研究需求。

建筑学

SwiftLLM的架构可以分为两个主要部分：控制平面和数据平面。

简单来说，控制面决定“计算什么”或“如何调度”，而数据面决定“如何计算”或“如何实现”并进行具体计算。它们以主从方式工作：控制平面充当主站，执行高级调度和协调，并将作业发送到数据平面，数据平面充当执行低级计算的工作线程。

控制平面的代码位于swiftllm/server目录中，包括Engine 、 Scheduler 、 API 服务器和TokenizationEngine等组件。数据平面的代码位于swiftllm/worker目录中，包括计算图的描述（在swiftllm/worker/model.py中）、模型中各层的实现（在swiftllm/layers中）以及 OpenAI Triton 内核（您可以将“内核”想象为在 GPU 上执行的函数（在swiftllm/kernels中）。

我们以玩具 API 服务器（位于swiftllm/server/api_server.py ）为例：

• 启动后，它使用EngineConfig创建Engine 。
• 之后，引擎通过Engine.initialize进行初始化，在其中创建Scheduler 、 TokenizationEngine和一组（由于不支持张量并行，目前只有一个）工作线程。然后它命令worker执行profile_num_blocks来计算GPU块的数量，之后引擎命令所有worker分配他们的KV缓存和KV交换。
• 最后，通过Engine.start_all_event_loops激活事件循环。在循环的每个步骤中，引擎向调度程序查询下一批要计算的请求，命令工作线程执行换入/换出，然后将批次发送给工作线程进行计算。
• API 服务器监听用户请求并与引擎交互来满足这些请求。

目前，控制平面 ( Engine ) 和数据平面 ( LlamaModel ) 驻留在同一节点上。实现张量并行/管道并行后，数据平面可能会分布到多个节点。

如何使用

我们提供两种使用 SwiftLLM 的方法：同时使用控制平面和数据平面，或仅使用数据平面。

如果您的想法足够简单或优雅，可以无缝集成到现有的控制平面中，那么您可以同时使用控制平面和数据平面。在另一种情况下，如果你想实现一个出色的IDE，你可能只能利用数据平面，并自己实现一个新的控制平面。

构建并运行

首先我们来搭建一下环境：

• 您可以从 Python >= 3.9 的干净 conda 环境开始，或者使用现有环境。如果您正在使用 conda 环境，请不要忘记激活它。
• 安装 PyTorch。注意根据您的硬件选择正确的版本。
• 通过pip install packaging安装packaging

然后就是安装：

• 通过git clone https://github.com/interestingLSY/swiftLLM.git克隆此存储库
• cd进入存储库 ( cd swiftLLM ) 并通过pip install -r requirements.txt安装其他依赖项。
• PyTorch 可能会为您安装 OpenAI Triton 的稳定版本。如果您想使用夜间版本以获得尖端性能但存在潜在问题，您可以卸载它并安装夜间版本。
• 运行pip install -e .将 SwiftLLM 安装到您的环境中。
• 通过pip install -e csrc安装一些 C 绑定

以下是一些示例：

• 目前 SwiftLLM 不支持自动从 HuggingFace 下载权重。您可能需要首先从 HuggingFace 克隆或下载模型权重。支持.bin格式和.safetensors格式。假设您的模型权重存储在/data/to/weight/ 。
• 对于离线服务示例，您可以尝试python3 examples/offline.py --model-path /data/to/weight 。该示例仅使用数据平面。如果您计划在没有控制平面的情况下使用 SwiftLLM，这是一个很好的起点。
• 对于使用Engine的在线服务示例，您可以尝试python3 examples/online.py --model-path /data/to/weight 。如果您计划同时使用控制平面和数据平面，这是一个很好的示例。
• 对于更复杂的示例，您可以参考swiftllm/server/api_server.py 。它启动了一个 API 服务器并提供了一个类似 vLLM 的在线服务界面。

表现

尽管很小（小也可以很可爱！），SwiftLLM 在性能上毫不妥协。我们在多种场景下评估了 SwiftLLM，并证明 SwiftLLM 可以实现与 vLLM 相当甚至更好的性能。

单次转发操作

第一个场景是“单个前向操作”，我们向模型提供一批输入，并让它生成一个输出令牌（相当于一个“前向”操作）。这是LLM推理（在线和离线）的基本操作，因此其性能至关重要。

这里我们使用 FP16 精度下配备 NVIDIA A100 80G PCIE / RTX 4090 GPU 的 LLaMA-3 7B 模型。结果如下所示（越低越好）：

可以看出，在相同设置下，SwiftLLM 可以达到与 vLLM 相当（甚至优于）的性能。

在线服务

第二个场景是“在线服务”，我们启动一个 API 服务器，从真实数据集中采样提示，并让模型生成补全。这是 LLM 在聊天机器人或代码完成等现实应用程序中使用的场景。

这里我们使用 ShareGPT 数据集对提示进行采样，并使用具有不同 lambda 的泊松过程来模拟不同的请求到达率。结果如下所示（越低越好）：

可以看到，在A100 80G PCIE上，SwiftLLM可以达到与vLLM相当的性能，而在RTX 4090上，SwiftLLM显着优于vLLM（主要是因为我们的控制平面的开销较低）。