🚀 弹性模型：DeepSeek - R1 - Distill - Llama - 8B。用于自助服务的最快、最灵活的模型

弹性模型是由TheStage AI的ANNA（自动化神经网络加速器）生成的模型。通过ANNA，你只需简单地滑动滑块，就能控制模型的大小、延迟和质量。对于每个模型，ANNA都会生成一系列优化模型：

• XL：数学上等效的神经网络，使用我们的DNN编译器进行优化。
• L：近乎无损的模型，在相应基准测试中的性能下降不到1%。
• M：更快的模型，准确率下降不到1.5%。
• S：最快的模型，准确率下降不到2%。

弹性模型的目标：

• 在推理时，为成本与质量的选择提供灵活性。
• 提供清晰的质量和延迟基准。
• 提供HF库（如transformers和diffusers）的接口，只需一行代码即可调用。
• 提供支持多种硬件的模型，这些模型已预编译，无需即时编译（JIT）。
• 为自托管提供最佳的模型和服务。

⚠️ 重要提示

需要注意的是，具体的质量下降情况会因模型而异。例如，S模型的准确率下降可能仅为0.5%。

🚀 快速开始

📦 安装指南

要使用我们的模型，只需在终端中运行以下命令：

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

# 目前我们需要你提供HF令牌
# 因为我们部分层使用了原始权重
# 同时也使用了模型配置
model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
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

# 推理过程与transformers库一样简单
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量化（使用int8数据类型），并使用了与ANNA相同的校准数据。S模型在速度上几乎相同，但质量要高得多，因为ANNA知道如何提高敏感层的量化质量！

质量基准

• MMLU：评估模型在包括科学、人文、工程等57个学科的通用知识，展示模型处理不同学术主题的能力。
• PIQA：通过关于日常物理交互的问题评估物理常识推理能力，展示模型对现实世界物理概念的理解。
• Arc Challenge：评估模型解决小学水平需要推理的选择题的能力，展示模型解决复杂推理任务的能力。
• Winogrande：通过句子完成任务评估常识推理能力，展示模型理解上下文和解决歧义的能力。

延迟基准

100输入/300输出；tok/s：

链接

• 平台：app.thestage.ai
• 订阅更新：TheStageAI X
• 联系邮箱：contact@thestage.ai

📄 许可证

本项目采用Apache - 2.0许可证。