Kimi悄悄开源了自家推理框架Mooncake~

Mooncake是为Kimi服务的平台,Kimi是由Moonshot AI提供的领先的大型语言模型(LLM)服务。

图片

Mooncake以KVCache为中心的解耦架构,将预填充和解码集群分开。它还利用GPU集群中未充分利用的CPU、DRAM和SSD资源来实现KVCache的解耦缓存。

图片

Mooncake的核心是其以KVCache为中心的调度器,它在最大化整体有效吞吐量的同时,满足与延迟相关的服务水平目标(SLOs)要求。与假设所有请求都将被处理的传统研究不同,Mooncake面临着由于高度过载场景带来的挑战。为了缓解这些问题,开发了一种基于预测的早期拒绝策略。实验表明,Mooncake在长上下文场景中表现出色。与基线方法相比,在某些模拟场景中,Mooncake可以实现高达525%的吞吐量增加,同时遵守SLOs。在实际工作负载下,Mooncake的创新架构使Kimi能够处理比之前多75%的请求。

图片

  • Mooncake的底层部分是传输引擎(Transfer Engine),它支持通过TCP、RDMA、基于NVIDIA GPUDirect的RDMA以及NVMe over Fabric(NVMe-of)协议进行快速、可靠和灵活的数据传输。与gloo(分布式PyTorch使用的)和TCP相比,Mooncake传输引擎具有最低的I/O延迟。

  • 基于传输引擎,实现了点对点存储库(P2P Store library),支持在集群中的节点之间共享临时对象(例如,检查点文件)。它避免了单台机器上的带宽饱和。

  • 此外,修改了vLLM,以便集成传输引擎。它通过利用RDMA设备,使预填充-解码解耦更加高效。

  • 未来,计划在传输引擎的基础上构建Mooncake Store,它支持池化的KVCache,以实现更灵活的预填充/解码(P/D)解耦。

https://arxiv.org/pdf/2407.00079https://github.com/kvcache-ai/Mooncake

### 集成 DeepSpeed 和 KIMI 框架 DeepSpeed 是一种用于训练大规模深度学习模型的库,能够显著提升分布式训练效率并优化资源利用。而 KIMI 可能是指特定于某个环境下的框架或工具集,在此假设其为与机器学习有关的技术栈的一部分。 为了在项目中集成这两个框架,可以遵循如下指南: #### 安装依赖项 首先安装必要的 Python 库来支持 DeepSpeed 功能: ```bash pip install deepspeed ``` 对于 KIMI 的具体安装方式取决于该组件的具体实现形式;如果是开源软件,则应参照官方文档获取最新的安装指令。 #### 修改 Web Configs 文件路径设置 当遇到启动错误提示无法找到 `PRODUCT_INFO_YAML_PATH` 时,应在项目的 `web_configs.py` 中加入指定配置文件的位置信息,并确保 `product_info.yaml` 已被放置到了相应的目录下[^1]。 #### 初始化 DeepSpeed 环境变量 为了让应用程序识别到 DeepSpeed 所需的各项参数,可以在运行脚本前设定一些环境变量,比如启用零冗余优化器 (ZeRO) 来节省内存消耗: ```bash export DEEPSPEED_ZERO_STAGE=2 ``` #### 创建兼容的 PyTorch 模型类 编写一个继承自 `torch.nn.Module` 的新类,以便后续可以通过调用 `.to('cuda')` 方法将其迁移到 GPU 上执行计算任务。同时定义好 forward 函数以描述数据流经网络的过程。 ```python import torch from transformers import BertForSequenceClassification class MyModel(torch.nn.Module): def __init__(self): super(MyModel, self).__init__() self.bert = BertForSequenceClassification.from_pretrained('bert-base-uncased') def forward(self, input_ids=None, attention_mask=None, labels=None): outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask, labels=labels) return outputs.logits ``` #### 构建 DeepSpeed 训练引擎实例 创建一个新的 Trainer 类用来管理整个训练过程中的细节操作,包括但不限于加载检查点、保存最佳权重等动作。这里会引入 DeepSpeed Engine 对象作为核心构件之一。 ```python import os import deepspeed def create_deepspeed_engine(model, optimizer, lr_scheduler, config_dict_or_path='ds_config.json'): model_engine, _, _, _ = deepspeed.initialize( args=None, model=model, model_parameters=[p for p in model.parameters() if p.requires_grad], config=config_dict_or_path) return model_engine ``` #### 调整 KVCache 大小策略 考虑到 Mooncake 架构的特点——即专注于改进KV缓存机制从而增强性能表现——可考虑采用多种手段缩小KV缓存量级,进而间接促进批处理规模扩大及命中率上升。这不仅有助于缓解硬件压力还能加快推理速度[^2]。 #### 实现示例代码片段 最后给出一段完整的Python程序样例供参考: ```python if __name__ == "__main__": from my_model_module import MyModel # 加载预训练好的BERT分类器 bert_classifier = MyModel() # 设置DeeepSpeed初始化所需的JSON配置文件位置 ds_config_file = './path/to/deepspeed/config' # 将模型转换为适合GPU加速的形式 device = 'cuda' if torch.cuda.is_available() else 'cpu' bert_classifier.to(device) # 获取优化器对象 optmizer = AdamW(bert_classifier.parameters(), lr=5e-5) # 合并上述要素构建最终版训练引擎 trainer = create_deepspeed_engine(bert_classifier, optmizer, None, ds_config_file) # 开始迭代更新权值... while not finished_training: ...
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值