tomeras1 commited on
Commit
a7db71f
1 Parent(s): 905f68a

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +161 -0
README.md CHANGED
@@ -1,3 +1,164 @@
1
  ---
 
2
  license: apache-2.0
 
 
 
 
 
3
  ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
+ library_name: transformers
3
  license: apache-2.0
4
+ tags:
5
+ - jamba
6
+ - mamba
7
+ - text-generation-inference
8
+ - moe
9
  ---
10
+
11
+ # Model Card for Jamba
12
+
13
+ Jamba is a state-of-the-art, hybrid SSM-Transformer LLM. It delivers throughput gains over traditional Transformer-based models, while outperforming or matching the leading models of its size class on most common benchmarks.
14
+
15
+ Jamba is the first production-scale Mamba implementation, which opens up interesting research and application opportunities. While this initial experimentation shows encouraging gains, we expect these to be further enhanced with future optimizations and explorations.
16
+
17
+ This model card is for the base version of Jamba. It’s a pretrained, mixture-of-experts (MoE) generative text model, with 12B active parameters and a total of 52B parameters across all experts. It supports a 256K context length, and can fit up to 140K tokens on a single 80GB GPU.
18
+
19
+ For full details of this model please read the [release blog post](https://www.ai21.com/blog/announcing-jamba).
20
+
21
+ ## Model Details
22
+
23
+ - **Developed by:** [AI21](https://www.ai21.com)
24
+ - **Model type:** Joint Attention and Mamba (Jamba)
25
+ - **License:** Apache 2.0
26
+ - **Context length:** 256K
27
+ - **Knowledge cutoff date:** March 5, 2024
28
+
29
+ ## Usage
30
+ ### Presequities
31
+ Jamba requires you use `transformers` version 4.38.2 or higher:
32
+ ```bash
33
+ pip install transformers>=4.38.2
34
+ ```
35
+
36
+ In order to run optimized Mamba implementations, you first need to install `mamba-ssm` and `causal-conv1d`:
37
+ ```bash
38
+ pip install mamba-ssm causal-conv1d>=1.2.0
39
+ ```
40
+ You also have to have the model on a CUDA device.
41
+
42
+ You can run the model not using the optimized Mamba kernels, but it is **not** recommended as it will result in significantly lower latencies. In order to do that, you'll need to specify `use_mamba_kernels=False` when loading the model.
43
+
44
+ ### Run the model
45
+ Please note that, at the moment, `trust_remote_code=True` is required for running the new Jamba architecture.
46
+ ```python
47
+ from transformers import AutoModelForCausalLM, AutoTokenizer
48
+
49
+ model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1",
50
+ trust_remote_code=True)
51
+ tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
52
+
53
+ input_ids = tokenizer("In the recent Super Bowl LVIII,", return_tensors='pt').to(model.device)["input_ids"]
54
+
55
+ outputs = model.generate(input_ids, max_new_tokens=6)
56
+
57
+ print(tokenizer.batch_decode(outputs))
58
+ # ["<|startoftext|>In the recent Super Bowl LVIII, the Kansas City Chiefs emerged victorious, defeating the San Francisco 49ers in a thrilling overtime showdown. The game was a nail-biter, with both teams showcasing their skills and determination.\n\nThe Chiefs, led by their star quarterback Patrick Mahomes, displayed their offensive prowess, while the 49ers, led by their strong defense, put up a tough fight. The game went into overtime, with the Chiefs ultimately securing the win with a touchdown.\n\nThe victory marked the Chiefs' second Super Bowl win in four years, solidifying their status as one of the top teams in the NFL. The game was a testament to the skill and talent of both teams, and a thrilling end to the NFL season.\n\nThe Super Bowl is not just about the game itself, but also about the halftime show and the commercials. This year's halftime show featured a star-studded lineup, including Usher, Alicia Keys, and Lil Jon. The show was a spectacle of music and dance, with the performers delivering an energetic and entertaining performance.\n"]
59
+ ```
60
+
61
+ <details>
62
+ <summary><strong>Loading the model in half precision</strong></summary>
63
+
64
+ The published checkpoint is saved in BF16. In order to load it into RAM in BF16/FP16, you need to specify `torch_dtype`:
65
+
66
+ ```python
67
+ from transformers import AutoModelForCausalLM
68
+ import torch
69
+ model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1",
70
+ trust_remote_code=True,
71
+ torch_dtype=torch.bfloat16) # you can also use torch_dtype=torch.float16
72
+ ```
73
+
74
+ When using half precision, you can enable the [FlashAttention2](https://github.com/Dao-AILab/flash-attention) implementation of the Attention blocks. In order to use it, you also need the model on a CUDA device. Since in this precision the model is to big to fit on a single 80GB GPU, you'll also need to parallelize it using [accelerate](https://huggingface.co/docs/accelerate/index):
75
+ ```python
76
+ from transformers import AutoModelForCausalLM
77
+ import torch
78
+ model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1",
79
+ trust_remote_code=True,
80
+ torch_dtype=torch.bfloat16,
81
+ attn_implementation="flash_attention_2",
82
+ device_map="auto")
83
+ ```
84
+
85
+ </details>
86
+ <details><summary><strong>Load the model in 8-bit</strong></summary>
87
+
88
+ **Using 8-bit precision, it is possible to fit up to 140K sequence lengths on a single 80GB GPU.** You can easily quantize the model to 8-bit using [bitsandbytes](https://huggingface.co/docs/bitsandbytes/index). In order to not degrade model quality, we recommend to exclude the Mamba blocks from the quantization:
89
+
90
+ ```python
91
+ from transformers import AutoModelForCausalLM, BitsAndBytesConfig
92
+ quantization_config = BitsAndBytesConfig(load_in_8bit=True,
93
+ llm_int8_skip_modules=["mamba"])
94
+ model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1",
95
+ trust_remote_code=True,
96
+ torch_dtype=torch.bfloat16,
97
+ attn_implementation="flash_attention_2",
98
+ quantization_config=quantization_config)
99
+ ```
100
+ </details>
101
+
102
+ ### Fine-tuning example
103
+ Jamba is a base model that can be fine-tuned for custom solutions (including for chat/instruct versions). You can fine-tune it using any technique of your choice. Here is an example of fine-tuning with the [PEFT](https://huggingface.co/docs/peft/index) library:
104
+
105
+ ```python
106
+ from datasets import load_dataset
107
+ from trl import SFTTrainer
108
+ from peft import LoraConfig
109
+ from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
110
+
111
+ tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
112
+ model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1", trust_remote_code=True, device_map='auto')
113
+
114
+ dataset = load_dataset("Abirate/english_quotes", split="train")
115
+ training_args = TrainingArguments(
116
+ output_dir="./results",
117
+ num_train_epochs=3,
118
+ per_device_train_batch_size=4,
119
+ logging_dir='./logs',
120
+ logging_steps=10,
121
+ learning_rate=2e-3
122
+ )
123
+ lora_config = LoraConfig(
124
+ r=8,
125
+ target_modules=["embed_tokens", "x_proj", "in_proj", "out_proj"],
126
+ task_type="CAUSAL_LM",
127
+ bias="none"
128
+ )
129
+ trainer = SFTTrainer(
130
+ model=model,
131
+ tokenizer=tokenizer,
132
+ args=training_args,
133
+ peft_config=lora_config,
134
+ train_dataset=dataset,
135
+ dataset_text_field="quote",
136
+ )
137
+
138
+ trainer.train()
139
+ ```
140
+
141
+ ## Results on common benchmarks
142
+ | Benchmark | Score |
143
+ |--------------|:-----:|
144
+ | HellaSwag | 87.1% |
145
+ | Arc Challenge | 64.4% |
146
+ | WinoGrande | 82.5% |
147
+ | PIQA | 83.2% |
148
+ | MMLU | 67.4% |
149
+ | BBH | 45.4% |
150
+ | TruthfulQA | 46.4% |
151
+ | GSM8K (CoT) | 59.9% |
152
+
153
+ It's crucial that the 'BOS' token is added to all prompts, which might not be enabled by default in all eval frameworks.
154
+
155
+
156
+ ## Notice
157
+ Jamba is a pretrained base model and did not undergo any alignment for instruct/chat interactions.
158
+
159
+ As a base model, Jamba is intended for use as a foundation layer for fine tuning, training, and developing custom solutions. Jamba does not have safety moderation mechanisms and guardrails should be added for responsible and safe use.
160
+
161
+ ## About AI21
162
+ AI21 builds reliable, practical, and scalable AI solutions for the enterprise.
163
+
164
+ Jamba is the first in AI21’s new family of models, and the Instruct version of Jamba is available in beta via the [AI21 platform](https://www.ai21.com/studio).