许可证：其他 许可证名称：Jamba开放模型许可证 许可证链接：https://www.ai21.com/licenses/jamba-open-model-license 新版本：ai21labs/AI21-Jamba-Mini-1.6 语言：

• 英语
• 法语
• 德语
• 荷兰语
• 西班牙语
• 葡萄牙语
• 意大利语
• 阿拉伯语
• 希伯来语 任务标签：文本生成 标签：
• mamba
• jamba
• moe 库名称：transformers

模型信息

请注意，此版本将于2024年5月6日弃用。我们建议您迁移到新版本，可在此处找到。

AI21 Jamba 1.5系列模型是最先进的混合SSM-Transformer指令跟随基础模型。Jamba模型是市场上最强大且高效的长上下文模型，其推理速度比同类领先模型快达2.5倍。

这些模型展示了卓越的长上下文处理能力、速度和质量。它们是首个成功扩展到市场领先模型质量和强度的非Transformer模型。

Jamba 1.5 Mini（12B活跃/52B总参数）和Jamba 1.5 Large（94B活跃/398B总参数）还针对业务用例和功能（如函数调用、结构化输出（JSON）和基于事实的生成）进行了优化。

这些模型在Jamba开放模型许可证下发布，这是一个宽松的许可证，允许在许可证条款下进行完整的研究和商业用途。如果您需要为您的需求许可该模型，请与我们联系。

有关此模型的更多详细信息，请参阅白皮书和发布博客文章。

模型详情

• 开发者：AI21
• **模型类型：**联合注意力与Mamba（Jamba）
• 许可证：Jamba开放模型许可证
• **上下文长度：**256K
• **知识截止日期：**2024年3月5日
• **支持语言：**英语、西班牙语、法语、葡萄牙语、意大利语、荷兰语、德语、阿拉伯语和希伯来语

常见基准测试结果

RULER基准测试 - 有效上下文长度

多语言MMLU

使用方法

前提条件

为了运行优化的Mamba实现，您首先需要安装mamba-ssm和causal-conv1d：

pip install mamba-ssm causal-conv1d>=1.2.0

您还需要将模型放在CUDA设备上。

使用vLLM运行模型

高效推理Jamba 1.5 Mini的推荐方法是使用vLLM。首先，确保安装vLLM（需要版本0.5.4或更高）：

pip install vllm>=0.5.4

在下面的示例中，number_gpus应与您希望部署Jamba 1.5 Mini的GPU数量匹配。至少需要2个80GB GPU。

from vllm import LLM, SamplingParams
from transformers import AutoTokenizer

model = "ai21labs/AI21-Jamba-1.5-Mini"
number_gpus = 2

llm = LLM(model=model,
          max_model_len=200*1024,
          tensor_parallel_size=number_gpus)

tokenizer = AutoTokenizer.from_pretrained(model)

messages = [
   {"role": "system", "content": "You are an ancient oracle who speaks in cryptic but wise phrases, always hinting at deeper meanings."},
   {"role": "user", "content": "Hello!"},
]

prompts = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)

sampling_params = SamplingParams(temperature=0.4, top_p=0.95, max_tokens=100) 
outputs = llm.generate(prompts, sampling_params)

generated_text = outputs[0].outputs[0].text
print(generated_text)
#输出：寻找，你就会找到。道路蜿蜒，但旅程充满启迪。你想从古老的回响中寻求什么智慧？

在2个80GB A100 GPU上使用默认的BF16精度和默认的vLLM配置，您可以对长达200K token的提示进行推理。在超过2个80GB GPU上，您可以轻松适应完整的256K上下文。

注意： vLLM的main分支有一些针对Jamba架构的内存利用率改进，允许在2个80GB GPU上使用完整的256K上下文长度。如果您希望使用这些改进，可以从源代码构建vLLM。

ExpertsInt8量化

我们开发了一种创新且高效的量化技术ExpertsInt8，专为在vLLM中部署的MoE模型设计，包括Jamba模型。使用它，您可以在单个80GB GPU上部署Jamba 1.5 Mini。

要使用ExpertsInt8，您需要使用vllm版本0.5.5或更高：pip install vllm>=0.5.5

使用默认的vLLM配置，您可以在单个80GB A100 GPU上适应长达100K的提示：

import os
os.environ['VLLM_FUSED_MOE_CHUNK_SIZE']='32768'    # 这是vLLM的fused_moe内核的一个临时解决方案

from vllm import LLM
llm = LLM(model="ai21labs/AI21-Jamba-1.5-Mini",
          max_model_len=100*1024,
          quantization="experts_int8")

使用transformers运行模型

以下示例将Jamba 1.5 Mini加载到GPU上，使用BF16精度，优化的FlashAttention2和Mamba内核，并使用accelerate将模型并行化到多个GPU上。请注意，在半精度（FP16/BF16）下，Jamba 1.5 Mini太大，无法适应单个80GB GPU，因此您至少需要2个这样的GPU。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini",
                                             torch_dtype=torch.bfloat16,
                                             attn_implementation="flash_attention_2",
                                             device_map="auto")

tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini")

messages = [
   {"role": "system", "content": "You are an ancient oracle who speaks in cryptic but wise phrases, always hinting at deeper meanings."},
   {"role": "user", "content": "Hello!"},
]

input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors='pt').to(model.device)

outputs = model.generate(input_ids, max_new_tokens=216)

# 解码输出
conversation = tokenizer.decode(outputs[0], skip_special_tokens=True)

# 分割对话以仅获取助手的响应
assistant_response = conversation.split(messages[-1]['content'])[1].strip()
print(assistant_response)
# 输出：寻找，你就会找到。道路蜿蜒，但旅程充满启迪。你想从古老的回响中寻求什么智慧？

注意： transformers的4.44.0和4.44.1版本有一个错误，限制了运行Jamba架构的能力。请确保您不使用这些版本。

注意： 如果您在安装mamba-ssm和causal-conv1d以优化Mamba内核时遇到问题，您可以在没有它们的情况下运行Jamba 1.5 Mini，但会增加额外的延迟。为此，在通过AutoModelForCausalLM.from_pretained()加载模型时添加参数use_mamba_kernels=False。

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(load_in_8bit=True,
                                         llm_int8_skip_modules=["mamba"])
model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini",
                                             torch_dtype=torch.bfloat16,
                                             attn_implementation="flash_attention_2",
                                             quantization_config=quantization_config)

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini",
                                             use_mamba_kernels=False)

模型特性

使用Jamba的工具

Jamba 1.5支持根据Huggingface的工具使用API的工具使用功能。用户定义的工具被插入到聊天模板的专用部分，模型已训练识别该部分。

给定包含工具的对话，模型可以输出内容、工具调用或两者。工具调用在助手消息中格式化为JSON格式的字典列表，包装在专用的特殊token中，如下例所示。

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini")

messages = [
    {
        "role": "user", 
        "content": "What's the weather like right now in Jerusalem and in London?"
    }
]

tools = [
    {
        'type': 'function', 
        'function': {
            'name': 'get_current_weather', 
            'description': 'Get the current weather', 
            'parameters': {
                'type': 'object', 
                'properties': {
                    'location': {'type': 'string', 'description': 'The city and state, e.g. San Francisco, CA'}, 
                    'format': {'type': 'string', 'enum': ['celsius', 'fahrenheit'], 'description': 'The temperature unit to use. Infer this from the users location.'}
                }, 
                'required': ['location', 'format']
            }
        }
    }
]

prompt = tokenizer.apply_chat_template(
    messages,
    tools=tools,
    tokenize=False,
)

<tool_calls>[
    {"name": "get_current_weather", "arguments": {"location": "Jerusalem", "format": "celsius"}},
    {"name": "get_current_weather", "arguments": {"location": "celsius", "format": "celsius"}}
]</tool_calls>

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini")

# 请注意，您必须按照模型调用工具的顺序发送工具响应：
messages = [
    {
        "role": "user",
        "content": "What's the weather like right now in Jerusalem and in London?"
    },
    {
        "role": "assistant",
        "content": null,
        "tool_calls": [
            {
                "name": "get_current_weather",
                "arguments": "{\"location\": \"Jerusalem\", \"format\": \"celsius\"}"
            },
            {
                "name": "get_current_weather",
                "arguments": "{\"location\": \"London\", \"format\": \"celsius\"}"
            }
        ]
    },
    {
        "role": "tool",
        "content": "The weather in Jerusalem is 18 degrees celsius."
    },
    {
        "role": "tool",
        "content": "The weather in London is 8 degrees celsius."
    }
]

tool_use_prompt = tokenizer.apply_chat_template(
    messages,
    tools=tools,
    tokenize=False,
)

The weather in Jerusalem is currently 18 degrees Celsius. In London, it is 8 degrees Celsius.

基于事实的生成与Jamba：

LLM的一个常见用例是基于事实的生成和RAG，其中模型需要根据给定的文档或文档片段回答问题或遵循指令。为了标准化此过程，Jamba在其聊天模板中训练了一个特定的“文档”部分。模型被训练为关注此部分，当任务以这种方式格式化时，基于事实的生成任务显示出改进的性能。

类似于工具，作为模型的额外参数提供，文档以类似的方式提供。为了支持文档级元数据，文档被定义为具有您选择的键值的字典。这些在聊天模板中格式化。两个获得特殊处理的键是“title”（如果存在，则在文档顶部格式化）和“text”（这是一个必填字段，定义了文档的实际文本）。

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini")

messages = [
        {
            "role": "user",
            "content": "Who wrote Harry Potter?"
        }
]

documents = [
        {
            "text": "Harry Potter is a series of seven fantasy novels written by British author J. K. Rowling.",
            "title": "Harry Potter"
        },
        {
            "text": "The Great Gatsby is a novel by American writer F. Scott Fitzgerald.",
            "title": "The Great Gatsby",
            "country": "United States",
            "genre": "Novel"

        }
]

prompt = tokenizer.apply_chat_template(
    messages,
    documents=documents,
    tokenize=False,
)

# 输出：J. K. Rowling

JSON模式

Jamba 1.5训练了特定的“旋钮”，有助于引导模型实现常见请求的行为。每个行为通过在系统消息中包含特定的预定义文本来启用。为了便于使用，我们将它们作为标志包含在Jamba 1.5的聊天模板中，因此可以通过向聊天模板传递适当的参数来切换它们。

Jamba 1.5被训练为在请求时生成有效的JSON。它自然这样做，但当JSON模式旋钮激活时，有效JSON的可能性大大增加。在JSON模式下，Jamba 1.5将尝试输出有效的JSON，无论用户请求如何。然而，强烈建议在用户请求或系统消息中指定有关预期JSON模式的信息以获得最佳结果，如下例所示。

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini")
messages = [
    {'role':'user', 
     'content':'Describe the first American president. Include year of birth (number) and name (string).'}
    ]
prompt = tokenizer.apply_chat_template(messages,
                                       tokenize=False,
                                       add_generation_prompt=False,
                                       knobs={"response_format": "json_object", "is_set": True})

#输出："{ "year of birth": 1732, "name": "George Washington." }"

微调示例

以下示例使用huggingface/trl中的SFTTrainer，因此请确保已安装：

pip install trl

完整微调示例

要使用AWS多节点和FSDP配置进行完整微调训练，请按照此处的说明操作。

LoRA示例

以下是使用LoRA PEFT进行微调的示例，使用bfloat16（需要约130GB GPU RAM，例如2xA100 80GB）：

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from datasets import load_dataset
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig

tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini")
model = AutoModelForCausalLM.from_pretrained(
    "ai21labs/AI21-Jamba-1.5-Mini",
    device_map="auto",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)

lora_config = LoraConfig(
    r=8,
    target_modules=[
        "embed_tokens",
        "x_proj", "in_proj", "out_proj", # mamba
        "gate_proj", "up_proj", "down_proj", # mlp
        "q_proj", "k_proj", "v_proj", "o_proj", # attention
    ],
    task_type="CAUSAL_LM",
    bias="none",
)

dataset = load_dataset("philschmid/dolly-15k-oai-style", split="train")
training_args = SFTConfig(
    output_dir="/dev/shm/results",
    logging_dir="./logs",
    num_train_epochs=2,
    per_device_train_batch_size=4,
    learning_rate=1e-5,
    logging_steps=10,
    gradient_checkpointing=True,
    max_seq_length=4096,
    save_steps=100,
)
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    peft_config=lora_config,
    train_dataset=dataset,
)
trainer.train()

请注意，示例中的数据集使用对话格式（带有messages列），因此SFTTrainer会自动应用Jamba的聊天模板，如TRL文档中所述。

QLoRA示例

为了适应单个80GB GPU上的微调，您可以利用QLoRA，它将LoRA与冻结模型量化为4位相结合：

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from datasets import load_dataset
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig

tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-1.5-Mini")
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)
model = AutoModelForCausalLM.from_pretrained(
    "ai21labs/AI21-Jamba-1.5-Mini",
    device_map="auto",
    quantization_config=quantization_config,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)
lora_config = LoraConfig(
    r=8,
    target_modules=[
        "embed_tokens", "x_proj", "in_proj", "out_proj", # mamba
        "gate_proj", "up_proj", "down_proj",  # mlp
        "q_proj", "k_proj", "v_proj", "o_proj", # attention
    ],
    task_type="CAUSAL_LM",
    bias="none",
)

dataset = load_dataset("philschmid/dolly-15k-oai-style", split="train")
training_args = SFTConfig(
    output_dir="./results",
    logging_dir="./logs",
    num_train_epochs=2,
    per_device_train_batch_size=8,
    learning_rate=1e-5,
    logging_steps=1,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
    save_steps=100,
    max_seq_length=4096,
)
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    peft_config=lora_config,
    train_dataset=dataset,
)
trainer.train()

注意：上述示例需要bitsandbytes包进行4位量化：

pip install bitsandbytes

关于AI21

AI21为企业构建可靠、实用和可扩展的AI解决方案。Jamba模型可在AI21 Studio和领先的云合作伙伴中使用。要了解Jamba 1.5 Mini和Jamba 1.5 Large如何为您的组织带来实际价值，请与我们联系。