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倍以上
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等前端
快速开始
环境准备
- 安装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界面。
注入式优化框架
KTransformers的核心是一个用户友好的、基于模板的注入框架。研究人员可以轻松地将原始torch模块替换为优化变体,并探索多种优化的协同效果。
使用示例
用户只需创建基于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技术的民主化和普及方面发挥重要作用。
