snowflake-arctic-instruct - 探索混合变换架构在企业AI中的高效应用

项目简介

Snowflake-Arctic-Instruct项目由Snowflake AI研究团队创建，是一种混合架构的transformer模型，采用密集-莫领域专家（MoE）的方式进行构建。Arctic模型分为基础版和指导调优版，目前已按照Apache-2.0许可证发布，允许自由使用这些模型进行研究、原型开发和产品应用。感兴趣的用户可以查看Snowflake官网博客获取更多信息，博客中包括了如何训练自定义MoE模型、生成高质量训练数据等多个方面的指南。

Snowflake还在其GitHub仓库中提供了详细的教程、代码示例等资料：

GitHub 项目地址

用户还可以通过Streamlit应用亲自体验模型的功能。

模型架构

Arctic结合了一个具有10B参数的密集transformer模型以及一个具备128x3.66B参量的MoE多层感知机，整体参数总数高达480B，但每次仅激活17B参数进行运算。通过这种Top-2门控方式，达到了计算效率与资源利用的平衡。关于模型架构的更多详细信息，如训练过程和数据使用情况，建议参考博客系列指南。

使用方法

目前，Arctic已支持通过transformers库进行使用，用户只需在调用AutoTokenizer和AutoModelForCausalLM时添加trust_remote_code=True即可。建议使用新版transformers库（版本4.39或以上）：

pip install transformers>=4.39.0

Arctic利用了一些DeepSpeed库的特性，因此需要安装DeepSpeed 0.14.2或更高版本：

pip install deepspeed>=0.14.2

推理示例

由于模型规模较大，建议使用像AWS p5.48xlarge或Azure ND96isr_H100_v5这样的高性能云实例来进行推理。在示例中，后台使用DeepSpeed提供的FP8量化，并可以通过指定q_bits=6实现FP6量化，同时对于max_memory参数设置为"150GiB"以支持FP量化。以下是使用实例代码：

import os
# 启用rapid下载功能
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from deepspeed.linear.config import QuantizationConfig

tokenizer = AutoTokenizer.from_pretrained(
    "Snowflake/snowflake-arctic-instruct",
    trust_remote_code=True
)
quant_config = QuantizationConfig(q_bits=8)

model = AutoModelForCausalLM.from_pretrained(
    "Snowflake/snowflake-arctic-instruct",
    trust_remote_code=True,
    low_cpu_mem_usage=True,
    device_map="auto",
    ds_quantization_config=quant_config,
    max_memory={i: "150GiB" for i in range(8)},
    torch_dtype=torch.bfloat16)

content = "5x + 35 = 7x - 60 + 10. Solve for x"
messages = [{"role": "user", "content": content}]
input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to("cuda")

outputs = model.generate(input_ids=input_ids, max_new_tokens=256)
print(tokenizer.decode(outputs[0]))

有关推理的更多代码示例，可前往Arctic的GitHub页面查阅：

纯HF库示例：点击查看
使用vLLM的教程：点击查看