KTransformers: 灵活高效的大语言模型推理优化框架

KTransformers简介

KTransformers是一个旨在增强🤗 Transformers体验的灵活框架,通过引入先进的内核优化和放置/并行策略,实现了大语言模型(LLM)推理的高效优化。作为一个以Python为中心的框架,KTransformers具有极强的可扩展性,用户只需一行代码即可实现优化模块的注入,就能获得与Transformers兼容的接口、符合OpenAI和Ollama标准的RESTful API,甚至是简化版的ChatGPT风格Web UI。

KTransformers的愿景是成为一个灵活的平台,用于尝试创新的LLM推理优化方法。它的核心优势在于能够在有限的计算资源下,实现大型语言模型的高效本地部署和运行。

主要特性与更新

KTransformers近期的主要更新包括:

支持在InternLM2.5-7B-Chat-1M模型下实现100万token上下文推理,仅需24GB显存和150GB内存。
将DeepseekV2所需显存从21GB降低到11GB。
支持llamafile作为线性后端。
支持多GPU使用。
新增对mixtral 87B和822B模型的支持。
在GPU上支持q2k、q3k、q5k反量化。
支持Windows原生运行。

showcase: 突破性能与效果

1. 桌面级100万token上下文推理

KTransformers在仅有24GB显存的桌面环境下,实现了对InternLM 2.5 7B模型的100万token上下文推理:

全精度bf16运算,仅需24GB显存和150GB内存
在100万token"大海捞针"测试中达到92.88%成功率
在128K token测试中达到100%成功率
生成速度达到16.91 tokens/s,比llama.cpp的全注意力方法快10倍以上

Single Needle Retrieval 128K

Single Needle Retrieval 1000K

2. 桌面级GPT-4水平的本地VSCode Copilot

KTransformers支持在仅有24GB显存的桌面环境下运行236B参数量的DeepSeek-Coder-V2模型:

使用Q4_K_M量化版本,仅需11GB显存和136GB内存
在BigCodeBench评测中得分超过GPT4-0613
通过MoE卸载和注入Llamafile、Marlin等先进内核,实现2K提示词预填充126 tokens/s,生成13.6 tokens/s的速度
封装为OpenAI和Ollama兼容的API,可无缝集成到Tabby等前端

DeepSeek-Coder-V2 Score

快速开始

环境准备

安装CUDA 12.1或更高版本
设置CUDA_HOME(Linux)或CUDA_PATH(Windows)环境变量
安装gcc、g++和cmake(Linux)
使用Conda创建Python 3.11虚拟环境
安装PyTorch、packaging和ninja

安装KTransformers

可以通过以下三种方式之一安装KTransformers:

使用Docker镜像

通过pip安装(Linux):

pip install ktransformers --no-build-isolation

从源码编译安装

本地聊天示例

以DeepSeek-V2-Lite-Chat-GGUF模型为例:

下载模型文件

运行本地聊天脚本:

python -m ktransformers.local_chat --model_path deepseek-ai/DeepSeek-V2-Lite-Chat --gguf_path ./DeepSeek-V2-Lite-Chat-GGUF

支持的模型

KTransformers目前支持多种大型语言模型,包括DeepSeek-V2、Qwen2-57B、Mixtral-8x7B等。每种模型都有相应的VRAM和DRAM需求,用户可根据自身硬件条件选择合适的模型。

RESTful API和Web UI

KTransformers提供了RESTful API服务和Web UI界面:

启动不带网页的服务:

ktransformers --model_path deepseek-ai/DeepSeek-V2-Lite-Chat --gguf_path /path/to/DeepSeek-V2-Lite-Chat-GGUF --port 10002

启动带网页的服务:

ktransformers --model_path deepseek-ai/DeepSeek-V2-Lite-Chat --gguf_path /path/to/DeepSeek-V2-Lite-Chat-GGUF  --port 10002 --web True

访问 http://localhost:10002/web/index.html#/chat 即可使用Web界面。

Web UI

注入式优化框架

KTransformers的核心是一个用户友好的、基于模板的注入框架。研究人员可以轻松地将原始torch模块替换为优化变体,并探索多种优化的协同效果。

Inject-Struction

使用示例

用户只需创建基于YAML的注入模板,并在使用Transformers模型之前调用optimize_and_load_gguf函数:

with torch.device("meta"):
    model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
optimize_and_load_gguf(model, optimize_rule_path, gguf_path, config)
...
generated = prefill_and_generate(model, tokenizer, input_tensor.cuda(), max_new_tokens=1000)

自定义模型优化

用户可以通过编写YAML模板来自定义模型优化规则。例如,将所有原始Linear模块替换为Marlin(一种先进的4位量化内核):

- match:
    name: "^model\.layers\...*$"
    class: torch.nn.Linear
  replace:
    class: ktransformers.operators.linear.KTransformerLinear
    device: "cpu"
    kwargs:
      generate_device: "cuda"
      generate_linear_type: "QuantizedLinearMarlin"

总结

KTransformers为大语言模型的本地部署和推理优化提供了一个灵活、高效的解决方案。它不仅支持多种先进模型,还提供了易于使用的API和Web界面。通过其独特的注入式优化框架,研究人员和开发者可以轻松尝试各种优化策略,推动LLM推理技术的进步。

KTransformers项目由清华大学MADSys小组和Approaching.AI的成员积极维护和开发。该项目欢迎新的贡献者加入,共同努力使KTransformers变得更快、更易用。随着大语言模型在各行各业的广泛应用,KTransformers无疑将在推动AI技术的民主化和普及方面发挥重要作用。