🚀 弹性模型:DeepSeek-R1-Distill-Qwen-7B
弹性模型是由TheStage AI的ANNA(自动化神经网络加速器)生成的模型,它能让你通过简单的滑动操作来控制模型的大小、延迟和质量。对于每个模型,ANNA都会生成一系列优化模型:
- XL:数学上等效的神经网络,使用我们的DNN编译器进行优化。
- L:近乎无损的模型,在相应基准测试中的性能下降不到1%。
- M:更快的模型,准确率下降不到1.5%。
- S:最快的模型,准确率下降不到2%。
弹性模型的目标
- 在推理时,为成本与质量的选择提供灵活性。
- 提供清晰的质量和延迟基准。
- 只需一行代码,即可提供HF库(transformers和diffusers)的接口。
- 提供支持广泛硬件的模型,这些模型已预编译,无需即时编译(JIT)。
- 为自建服务提供最佳的模型和服务。
⚠️ 重要提示
具体的质量下降情况可能因模型而异。例如,S模型的性能下降可能只有0.5%。

🚀 快速开始
✨ 主要特性
弹性模型由TheStage AI的ANNA自动化神经网络加速器生成,具备以下特性:
- 可通过简单操作灵活控制模型大小、延迟和质量。
- 针对不同性能需求,提供多种优化模型版本。
- 提供清晰的质量和延迟基准,方便用户选择。
- 支持与HF库无缝对接,使用便捷。
- 预编译模型,支持广泛硬件,无需即时编译。
📦 安装指南
要使用我们的模型,只需在终端中运行以下命令:
pip install thestage
pip install elastic_models[nvidia]\
--index-url https://thestage.jfrog.io/artifactory/api/pypi/pypi-thestage-ai-production/simple\
--extra-index-url https://pypi.nvidia.com\
--extra-index-url https://pypi.org/simple
pip install flash_attn==2.7.3 --no-build-isolation
pip uninstall apex
然后访问 app.thestage.ai 登录,并从个人资料页面生成API令牌。设置API令牌的命令如下:
thestage config set --api-token <YOUR_API_TOKEN>
恭喜,现在你可以使用加速模型了!
💻 使用示例
基础用法
要推理我们的模型,你只需将 transformers
导入替换为 elastic_models.transformers
:
import torch
from transformers import AutoTokenizer
from elastic_models.transformers import AutoModelForCausalLM
model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
hf_token = ''
device = torch.device("cuda")
tokenizer = AutoTokenizer.from_pretrained(
model_name, token=hf_token
)
model = AutoModelForCausalLM.from_pretrained(
model_name,
token=hf_token,
torch_dtype=torch.bfloat16,
attn_implementation="sdpa",
mode='S'
).to(device)
model.generation_config.pad_token_id = tokenizer.eos_token_id
prompt = "Describe basics of DNNs quantization."
messages = [
{
"role": "system",
"content": "You are a search bot, answer on user text queries."
},
{
"role": "user",
"content": prompt
}
]
chat_prompt = tokenizer.apply_chat_template(
messages, add_generation_prompt=True, tokenize=False
)
inputs = tokenizer(chat_prompt, return_tensors="pt")
inputs.to(device)
with torch.inference_mode():
generate_ids = model.generate(**inputs, max_length=500)
input_len = inputs['input_ids'].shape[1]
generate_ids = generate_ids[:, input_len:]
output = tokenizer.batch_decode(
generate_ids,
skip_special_tokens=True,
clean_up_tokenization_spaces=False
)[0]
print(f"# Q:\n{prompt}\n")
print(f"# A:\n{output}\n")
系统要求
- GPU:H100、L40s
- CPU:AMD、Intel
- Python:3.10 - 3.12
📚 详细文档
基准测试
基准测试是模型加速过程中最重要的步骤之一。我们旨在为使用我们算法的模型提供清晰的性能指标。W8A8, int8
列表示我们对所有线性层应用了W8A8量化,并使用了与ANNA相同的校准数据。S模型在实现几乎相同速度的同时,具有更高的质量,因为ANNA知道如何提高敏感层的量化质量!
质量基准
指标/模型 |
S |
M |
L |
XL |
原始模型 |
W8A8, int8 |
arc_challenge |
41.00 |
40.90 |
42.50 |
42.20 |
42.20 |
21.00 |
mmlu |
52.00 |
53.80 |
55.10 |
55.20 |
55.20 |
48.00 |
piqa |
68.40 |
70.60 |
70.70 |
70.50 |
70.50 |
50.40 |
winogrande |
60.10 |
60.90 |
60.20 |
60.10 |
60.10 |
51.10 |
- MMLU:评估涵盖科学、人文、工程等57个学科的通用知识,展示模型处理多样化学术主题的能力。
- PIQA:通过关于日常物理交互的问题评估物理常识推理能力,展示模型对现实世界物理概念的理解。
- Arc Challenge:评估需要推理的小学水平多项选择题,展示模型解决复杂推理任务的能力。
- Winogrande:通过句子完成任务评估常识推理能力,展示模型理解上下文和解决歧义的能力。
延迟基准
100输入/300输出;tok/s:
GPU/模型 |
S |
M |
L |
XL |
原始模型 |
W8A8, int8 |
H100 |
204 |
185 |
161 |
135 |
62 |
205 |
L40S |
78 |
69 |
60 |
47 |
43 |
78 |
🔗 链接
📄 许可证
本项目采用Apache-2.0许可证。