🚀 彈性模型:Mistral-7B-Instruct-v0.3
彈性模型是由TheStage AI的ANNA(自動化神經網絡加速器)生成的模型。ANNA允許你通過簡單的滑動操作來控制模型的大小、延遲和質量。對於每個模型,ANNA都會生成一系列優化模型:
- XL:數學上等效的神經網絡,使用我們的DNN編譯器進行優化。
- L:近乎無損的模型,在相應基準測試中的性能下降小於1%。
- M:更快的模型,準確率下降小於1.5%。
- S:最快的模型,準確率下降小於2%。
彈性模型的目標:
- 在推理時提供成本與質量選擇的靈活性。
- 提供清晰的質量和延遲基準。
- 提供HF庫(transformers和diffusers)的接口,只需一行代碼。
- 提供支持廣泛硬件的模型,這些模型是預編譯的,無需JIT。
- 為自託管提供最佳的模型和服務。
⚠️ 重要提示
具體的質量下降可能因模型而異。例如,S模型的性能下降也可能只有0.5%。

🚀 快速開始
✨ 主要特性
- 可通過ANNA靈活控制模型大小、延遲和質量。
- 提供多種優化模型版本(XL、L、M、S)。
- 提供清晰的質量和延遲基準。
- 支持廣泛硬件,預編譯無需JIT。
- 提供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
- 生成並設置API令牌:
前往 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 = "mistralai/Mistral-7B-Instruct-v0.3"
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量化(使用int8數據類型),並使用了與ANNA相同的校準數據。S模型實現了幾乎相同的速度,但質量更高,因為ANNA知道如何提高敏感層的量化質量!
質量基準
指標/模型 |
S |
M |
L |
XL |
原始模型 |
W8A8, int8 |
MMLU |
59.7 |
60.1 |
60.8 |
61.4 |
61.4 |
28 |
PIQA |
80.8 |
82 |
81.7 |
81.5 |
81.5 |
65.3 |
Arc Challenge |
56.6 |
55.1 |
56.8 |
57.4 |
57.4 |
33.2 |
Winogrande |
73.2 |
72.3 |
73.2 |
74.1 |
74.1 |
57 |
- MMLU:評估跨57個學科(包括科學、人文、工程等)的通用知識,展示模型處理多樣化學術主題的能力。
- PIQA:通過關於日常物理交互的問題評估物理常識推理能力,展示模型對現實世界物理概念的理解。
- Arc Challenge:評估需要推理的小學水平多項選擇題,展示模型解決複雜推理任務的能力。
- Winogrande:通過句子完成任務評估常識推理能力,展示模型理解上下文和解決歧義的能力。
延遲基準
100輸入/300輸出;tok/s:
GPU/模型 |
S |
M |
L |
XL |
原始模型 |
W8A8, int8 |
H100 |
186 |
180 |
168 |
136 |
48 |
192 |
L40s |
79 |
68 |
59 |
47 |
38 |
82 |
📄 許可證
本項目採用Apache-2.0許可證。
🔗 相關鏈接