Snowflake/snowflake-arctic-base

1---2license: apache-2.03tags:4- snowflake5- arctic6- moe7---8 9## Model Details10 11Arctic is a dense-MoE Hybrid transformer architecture pre-trained from scratch by the Snowflake AI 12Research Team. We are releasing model checkpoints for both the base and instruct-tuned versions of 13Arctic under an Apache-2.0 license. This means you can use them freely in your own research, 14prototypes, and products. Please see our blog 15[Snowflake Arctic: The Best LLM for Enterprise AI — Efficiently Intelligent, Truly Open](https://www.snowflake.com/blog/arctic-open-efficient-foundation-language-models-snowflake/) 16for more information on Arctic and links to other relevant resources such as our series of cookbooks 17covering topics around training your own custom MoE models, how to produce high-quality training data, 18and much more.19 20* [Arctic-Base](https://huggingface.co/Snowflake/snowflake-arctic-base/)21* [Arctic-Instruct](https://huggingface.co/Snowflake/snowflake-arctic-instruct/)22 23For the latest details about Snowflake Arctic including tutorials, etc. please refer to our github repo: 24* https://github.com/Snowflake-Labs/snowflake-arctic25 26**Model developers** Snowflake AI Research Team27 28**License** Apache-2.029 30**Input** Models input text only.31 32**Output** Models generate text and code only.33 34**Model Release Date** April, 24th 2024.35 36## Model Architecture37 38Arctic combines a 10B dense transformer model with a residual 128x3.66B MoE MLP resulting in 480B 39total and 17B active parameters chosen using a top-2 gating. For more details about Arctic's model40Architecture, training process, data, etc. [see our series of cookbooks](https://www.snowflake.com/en/data-cloud/arctic/cookbook/).41 42## Usage43 44Arctic is currently supported with `transformers` by leveraging the 45[custom code feature](https://huggingface.co/docs/transformers/en/custom_models#using-a-model-with-custom-code), 46to use this you simply need to add `trust_remote_code=True` to your AutoTokenizer and AutoModelForCausalLM calls.47However, we recommend that you use a `transformers` version at or above 4.39:48 49```python50pip install transformers>=4.39.051```52 53Arctic leverages several features from [DeepSpeed](https://github.com/microsoft/DeepSpeed), you will need to 54install the DeepSpeed 0.14.2 or higher to get all of these required features:55 56```python57pip install deepspeed>=0.14.258```59 60### Inference examples61 62Due to the model size we recommend using a single 8xH100 instance from your63favorite cloud provider such as: AWS [p5.48xlarge](https://aws.amazon.com/ec2/instance-types/p5/), 64Azure [ND96isr_H100_v5](https://learn.microsoft.com/en-us/azure/virtual-machines/nd-h100-v5-series), etc.65 66In this example we are using FP8 quantization provided by DeepSpeed in the backend, we can also use FP6 67quantization by specifying `q_bits=6` in the `QuantizationConfig` config. The `"150GiB"` setting 68for max_memory is required until we can get DeepSpeed's FP quantization supported natively as a69[HFQuantizer](https://huggingface.co/docs/transformers/main/en/hf_quantizer#build-a-new-hfquantizer-class) which we 70are actively working on.71 72```python73import os74# enable hf_transfer for faster ckpt download75os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"76 77import torch78from transformers import AutoModelForCausalLM, AutoTokenizer79from deepspeed.linear.config import QuantizationConfig80 81tokenizer = AutoTokenizer.from_pretrained(82    "Snowflake/snowflake-arctic-instruct",83    trust_remote_code=True84)85quant_config = QuantizationConfig(q_bits=8)86 87model = AutoModelForCausalLM.from_pretrained(88    "Snowflake/snowflake-arctic-instruct",89    trust_remote_code=True,90    low_cpu_mem_usage=True,91    device_map="auto",92    ds_quantization_config=quant_config,93    max_memory={i: "150GiB" for i in range(8)},94    torch_dtype=torch.bfloat16)95 96 97content = "5x + 35 = 7x - 60 + 10. Solve for x"98messages = [{"role": "user", "content": content}]99input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to("cuda")100 101outputs = model.generate(input_ids=input_ids, max_new_tokens=256)102print(tokenizer.decode(outputs[0]))103```104 105The Arctic github page has additional code snippets and examples around running inference:106 107* Example with pure-HF: https://github.com/Snowflake-Labs/snowflake-arctic/blob/main/inference108* Tutorial using vLLM: https://github.com/Snowflake-Labs/snowflake-arctic/tree/main/inference/vllm