长亭百川云 - 文章详情

KTransformers:一个灵活的框架,用于体验最先进的 LLM 推理优化

CTOLib码库

186

2024-07-29

KTransformers 是一个灵活的、以 Python 为中心的框架,其核心设计理念是可扩展性。通过一行代码实现和注入一个优化的模块,用户可以获得与 Transformers 兼容的接口、符合 OpenAI 和 Ollama 的 RESTful API,甚至简化的 ChatGPT 式 Web UI。

我们希望 KTransformers 成为一个灵活的平台,用于试验创新的 LLM 推理优化。如果您需要其他功能,请告诉我们。

请注意,在使用 DeepSeek 和 QWen 时,您需要遵守其相应的模型许可。

安装 KTransformers:

•使用源代码安装:

pip install -r requirements-local_chat.txt
  pip install . --no-build-isolation

启动 KTransformers:

•无网站启动:

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

•使用 Transformers 启动服务器:

ktransformers --type transformers --model_path /mnt/data/model/Qwen2-0.5B-Instruct --port 10002 --web True

访问网站:

•使用 URL http://localhost:10002/web/index.html#/chat
 访问网站。

RESTful API 服务器:

•更多关于 RESTful API 服务器的信息,请点击这里。•集成 Tabby 的示例,请点击这里。

注入教程:

KTransformers 的核心是一个用户友好的、基于模板的注入框架。这允许研究人员轻松地用优化后的变体替换原始的 torch 模块。它还简化了组合多个优化的过程,允许探索它们的协同效应。

鉴于 vLLM 已经成为一个用于大规模部署优化的优秀框架,KTransformers 特别关注受有限资源限制的本地部署。我们特别关注异构计算机会,例如量化模型的 GPU/CPU 卸载。例如,我们支持 Llamafile 和 Marlin 内核,分别用于 CPU 和 GPU。更多细节,请点击这里。

使用提供的内核:

•用户只需要创建一个基于 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 模板示例:

•下面是一个使用 Marlin(一个先进的 4 位量化内核)替换所有原始 Linear 模块的 YAML 模板示例。

- 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/optimize/optimize_rules
 目录中找到用于优化 DeepSeek-V2 和 Qwen2-57B-A14(两个 SOTA MoE 模型)的示例规则模板。这些模板用于为 local_chat.py
 演示提供支持。

使用 DeepSeek-V2 作为示例的注入详细说明,请点击这里。

致谢和贡献者:

KTransformer 的开发基于 Transformers 提供的灵活且通用的框架。我们也受益于 GGUF/GGML、Llamafile 和 Marlin 等先进内核。我们计划通过上游我们的修改来回馈社区。

KTransformer 由清华大学 MADSys 团队的贡献者和 Approaching.AI 的成员积极维护和开发。我们欢迎新的贡献者加入我们,共同努力使 KTransformer 更快、更易于使用。

相关推荐
关注或联系我们
添加百川云公众号,移动管理云安全产品
咨询热线:
4000-327-707
百川公众号
百川公众号
百川云客服
百川云客服

Copyright ©2024 北京长亭科技有限公司
icon
京ICP备 2024055124号-2