Elastic DeepSeek R1 Distill Llama 8B

🚀 彈性模型：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許可證。