Add support to ArcticForCausalLM

[maziyarpanahi]

First open LLM from @SnowflakeDB! Arctic is 480B Dense-MoE with a 10B dense transformer model and a 128x3.66B MoE MLP designed specifically for enterprise AI. 🤔

TL;DR:
🧠 480B parameters with 17B active during generation
👨‍🏫 128 experts with 2 active in generation
2️⃣ Instruct & Base versions released
🏙️ Focused on Enterprise task (Code, SQL, Reasoning, Following)
🔓 Released under Apache 2.0
🗻 in fp16 ~900GB Memory & in int4 ~240GB
🤗 Available on @huggingface

🏋🏻 Trained with DeepSpeed-MoE

Blog: https://snowflake.com/blog/arctic-open-efficient-foundation-language-models-snowflake/

Models: https://huggingface.co/Snowflake/snowflake-arctic-instruct

[fairydreaming]

I looked briefly into possible support of this model in convert-hf-to-gguf.py and found that there are some tensors causing Can not map tensor errors: model.layers.0.residual_layernorm.weight, model.layers.0.residual_mlp.w1.weight. I checked the Snowflake transformers github arctic branch and found that they can be likely ignored, since they are created only when:

        if self.parallel_attn_mlp_res:
            self.residual_layernorm = ArcticRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
            self.residual_mlp =  ArcticMLP(config,

and the default value of parallel_attn_mlp_res is False (see src/transformers/models/arctic/configuration_arctic.py file).

It's possible that only the following changes are needed to make convert-hf-to-gguf.py work:

diff --git a/convert-hf-to-gguf.py b/convert-hf-to-gguf.py
index 5763b666..d02f7aea 100755
--- a/convert-hf-to-gguf.py
+++ b/convert-hf-to-gguf.py
@@ -1303,7 +1303,7 @@ class StableLMModel(Model):
             self.gguf_writer.add_tensor(new_name, data)
 
 
-@Model.register("LlamaForCausalLM", "MistralForCausalLM", "MixtralForCausalLM")
+@Model.register("LlamaForCausalLM", "MistralForCausalLM", "MixtralForCausalLM", "ArcticForCausalLM")
 class LlamaModel(Model):
     model_arch = gguf.MODEL_ARCH.LLAMA
 
@@ -1345,7 +1345,9 @@ class LlamaModel(Model):
         experts = dict()
         for name, data_torch in self.get_tensors():
             # we don't need these
-            if name.endswith((".attention.masked_bias", ".attention.bias", ".attention.rotary_emb.inv_freq")):
+            if name.endswith((".attention.masked_bias", ".attention.bias", ".attention.rotary_emb.inv_freq", ".residual_layernorm.weight")):
+                continue
+            if re.match(r".*\.residual_mlp\.w\d+\.weight", name):
                 continue
 
             old_dtype = data_torch.dtype

So a brave soul may try to apply this patch and make the quants, but someone with more expertise shall verify this. I ran it on Amazon EC2 instance and it worked up to writing the GGUF file. Good luck!

[christianazinn]

I'm currently downloading the model weights (Instruct) and will see if conversion works on my local machine (don't have access to cloud compute), but it'd take a while. If someone with cloud computing capabilities could take a go, that would be way faster, but I'll try in the meantime.

For what it's worth, it uses the ChatML prompt template and LLaMA 2 tokenizer. I don't know anything else about the architecture yet.

[fairydreaming]

Based on the image in this article it's a hybrid combining a dense transformer with a residual MoE component. I guess a new LLM architecture will be necessary to support this in llama.cpp.

[maziyarpanahi]

The team behind it (very capable people!) are willing to help making this model more accessible. I have seen their PRs in DeepSpeed to introduce new FP [6,8,12]:

• FP [6,8,12] quantizer op deepspeedai/DeepSpeed#5336
• Optimize the fp-dequantizer to get high memory-BW utilization deepspeedai/DeepSpeed#5373
• Selective dequantization deepspeedai/DeepSpeed#5375

If needed, we can reach out to them for some insight into how to go about this better?

[sorasoras]

The team behind it (very capable people!) are willing to help making this model more accessible. I have seen their PRs in DeepSpeed to introduce new FP [6,8,12]:

• FP [6,8,12] quantizer op deepspeedai/DeepSpeed#5336
• Optimize the fp-dequantizer to get high memory-BW utilization deepspeedai/DeepSpeed#5373
• Selective dequantization deepspeedai/DeepSpeed#5375

If needed, we can reach out to them for some insight into how to go about this better?

I was thinking only load dense part onto GPU and leave moe part on CPU ram. That way a lot of people can run this in decent speed

[maziyarpanahi]

That's actually not a bad idea! Would it be possible to make this more dynamic in all MoE models, allowing users to select whether the experts to be offloaded on GPUs/CPUs?

[christianazinn]

Not an expert (haha) but I don't know if there's any way to determine which layers correspond to which experts.

In the first place, the way MoE inference works shouldn't really allow this to be possible - the gate network sends different tokens to different experts, so you'd need all those experts to be on GPU for any increase in inference speed. I don't know how it works with the additional dense part in this particular architecture, but I assume it would get bottlenecked by the experts; I would be happy to be wrong, though.

[cpumaxx]

I was able to convert and quant this model to Q8 with convert.py --skip-unknown, but needed to increase LLAMA_MAX_EXPERTS from 60 to at least 128 in llama.cpp to get it to load the model.
The output, however, was only partially coherent.
I have attached the output from loading with current up-to-date llama.cpp for reference

[christianazinn]

I was able to convert and quant this model to Q8 with convert.py --skip-unknown, but needed to increase LLAMA_MAX_EXPERTS from 60 to at least 128 in llama.cpp to get it to load the model. The output, however, was only partially coherent. I have attached the output from loading with current up-to-date llama.cpp for reference

Can you post some examples of what you mean by "partially coherent," and might conversion with LLAMA_MAX_EXPERTS=60 have affected it? And how large is the Q8_0 file?

[cpumaxx]

might conversion with LLAMA_MAX_EXPERTS=60 have affected it?
The conversion script didn't need any modifications to run, although whether it worked as intended is unknown.
how large is the Q8_0 file?
Q8 is 472G.
Can you post some examples of what you mean by "partially coherent,"
Response to a simple -p "Hello" is below. "Partially coherent" might have been giving it too much credit:

[christianazinn]

might conversion with LLAMA_MAX_EXPERTS=60 have affected it?
The conversion script didn't need any modifications to run, although whether it worked as intended is unknown.
how large is the Q8_0 file?
Q8 is 472G.
Can you post some examples of what you mean by "partially coherent,"
Response to a simple -p "Hello" is below. "Partially coherent" might have been giving it too much credit:

Thanks. --skip-unknown probably did the same as @fairydreaming's modification, and hopefully it defaulted to assuming it was a Llama architecture model. It seems like it at least generates tokens properly, which is a plus! :P convert-hf-to-gguf.py with the suggested edits is underway on my end to f16, and I'll see if I get the same results. It's likely an issue with the hybrid architecture.

[fairydreaming]

I was able to convert and quant this model to Q8 with convert.py --skip-unknown, but needed to increase LLAMA_MAX_EXPERTS from 60 to at least 128 in llama.cpp to get it to load the model. The output, however, was only partially coherent. I have attached the output from loading with current up-to-date llama.cpp for reference

I wonder why does it display BOS and EOS tokens as some weird characters instead of <|im_start|> and <|im_end|>.

[fairydreaming]

Some quants for testing: https://huggingface.co/sszymczyk/snowflake-arctic-instruct-GGUF/

[fairydreaming]

I think ignoring these tensors likely won't work after all, I just found that in https://huggingface.co/Snowflake/snowflake-arctic-instruct/blob/main/config.json they have parallel_attn_mlp_res set to true, so residual_layernorm and residual_mlp are created and used. Sorry for the confusion.