Add support for DBRX models: dbrx-base and dbrx-instruct

[maziyarpanahi]

Prerequisites

Please answer the following questions for yourself before submitting an issue.

• I am running the latest code. Development is very rapid so there are no tagged versions as of now.
• I carefully followed the README.md.
• I searched using keywords relevant to my issue to make sure that I am creating a new issue that is not already open (or closed).
• I reviewed the Discussions, and have a new bug or useful enhancement to share.

Feature Description

Databricks just released 2 new models called DBRX (base and instruct). They have their own architecture:

{
  "architectures": [
    "DbrxForCausalLM"
  ],
  "attn_config": {
    "clip_qkv": 8,
    "kv_n_heads": 8,
    "model_type": "",
    "rope_theta": 500000
  },
  "auto_map": {
    "AutoConfig": "configuration_dbrx.DbrxConfig",
    "AutoModelForCausalLM": "modeling_dbrx.DbrxForCausalLM"
  },
  "d_model": 6144,
  "emb_pdrop": 0.0,
  "ffn_config": {
    "ffn_hidden_size": 10752,
    "model_type": "",
    "moe_jitter_eps": 0,
    "moe_loss_weight": 0.05,
    "moe_num_experts": 16,
    "moe_top_k": 4
  },
  "initializer_range": 0.02,
  "max_seq_len": 32768,
  "model_type": "dbrx",
  "n_heads": 48,
  "n_layers": 40,
  "output_router_logits": false,
  "resid_pdrop": 0.0,
  "router_aux_loss_coef": 0.05,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.38.2",
  "use_cache": true,
  "vocab_size": 100352
}

Motivation

These models are superior to the predecessors like Llama-2 or Mixtral (even though they are larger), the community can really benefit from these two and the fine-tuned models that come after.

https://huggingface.co/databricks/dbrx-instruct

Possible Implementation

If you have an idea as to how it can be implemented, please write a detailed description. Feel free to give links to external sources or share visuals that might be helpful to understand the details better.

python llama.cpp/convert-hf-to-gguf.py

Traceback (most recent call last):
  File "/llama.cpp/convert-hf-to-gguf.py", line 2099, in <module>
    main()
  File "/llama.cpp/convert-hf-to-gguf.py", line 2079, in main
    model_class = Model.from_model_architecture(hparams["architectures"][0])
  File "/llama.cpp/convert-hf-to-gguf.py", line 215, in from_model_architecture
    raise NotImplementedError(f'Architecture {arch!r} not supported!') from None
NotImplementedError: Architecture 'DbrxForCausalLM' not supported!

python llama.cpp/convert.py

  File "/llama.cpp/convert.py", line 1486, in <module>
    main()
  File "/llama.cpp/convert.py", line 1422, in main
    model_plus = load_some_model(args.model)
  File "/llama.cpp/convert.py", line 1291, in load_some_model
    model_plus = merge_multifile_models(models_plus)
  File "/llama.cpp/convert.py", line 747, in merge_multifile_models
    model = merge_sharded([mp.model for mp in models_plus])
  File "/llama.cpp/convert.py", line 726, in merge_sharded
    return {name: convert(name) for name in names}
  File "/llama.cpp/convert.py", line 726, in <dictcomp>
    return {name: convert(name) for name in names}
  File "/llama.cpp/convert.py", line 701, in convert
    lazy_tensors: list[LazyTensor] = [model[name] for model in models]
  File "/llama.cpp/convert.py", line 701, in <listcomp>
    lazy_tensors: list[LazyTensor] = [model[name] for model in models]
KeyError: 'transformer.blocks.0.ffn.experts.mlp.w1'

Dbrx is a mixture-of-experts model, which each FFN is divided into 16 experts and only 4 are activated at any given time. We build on MegaBlocks
https://github.com/databricks/megablocks

[abhi-mosaic]

Hi! DBRX researcher here, happy to help out however I can!

The architecture is quite similar to Mixtral, which is already supported in this framework. The modeling source code for DBRX is available on the HF Hub here: https://huggingface.co/databricks/dbrx-instruct/blob/main/modeling_dbrx.py

The main differences vs. Mixtral as far as I can tell:

• QKV clipping: https://huggingface.co/databricks/dbrx-instruct/blob/464e701f50aef4c1b59c81fb5667819a5d08e108/modeling_dbrx.py#L317-L318
• LayerNorm rather than RMSNorm: https://huggingface.co/databricks/dbrx-instruct/blob/464e701f50aef4c1b59c81fb5667819a5d08e108/modeling_dbrx.py#L642
• the weights of the DbrxExpertGLU layer are fused along the experts , aka transformer.blocks.12.ffn.experts.mlp.w1, not transformer.blocks.12.ffn.experts.mlp.list_w1.0.weight. This is what you will find in the .safetensors weight files, you can see the mapping explicitly here in the model.safetensors.index.json: https://huggingface.co/databricks/dbrx-instruct/blob/main/model.safetensors.index.json
  • So if the llama.cpp weight loader expects to load individual tensors for each expert, I think there needs to be some kind of sharding/renaming of the fused tensors (e.g. ...mlp.w1) after they are loaded from disk.
  • DbrxExpertGLU: https://huggingface.co/databricks/dbrx-instruct/blob/464e701f50aef4c1b59c81fb5667819a5d08e108/modeling_dbrx.py#L749

Please let me know if you have any questions!

[abhi-mosaic]

The model is ~132B params so I think the expected memory usage is:

• 264GB with 16-bit, already possible using our HF quickstart code: https://huggingface.co/databricks/dbrx-instruct#quickstart-guide
• expected 132GB with 8-bit
• expected 99GB with 6-bit
• expected 66GB with 4-bit

[ggerganov]

@abhi-mosaic Thanks for the pointers. We do split the experts in separate tensors at the moment, but it is something that we planned to change: #6082

Seems like now is the time do that

[maziyarpanahi]

Thanks @abhi-mosaic for all the complete and detailed explanations.

@ggerganov I have a big server, I can test any PR from 16bit all the way down to 2bit. (I have the model already downloaded and ready)

[moshemalawach]

Same, got big servers with very fast and plentyful ram channels, so can try on CPU all the sizes.

[sirus20x6]

Put me in coach, I'm ready to play, today.

[veryvanya]

happy to test on my server

[simsim314]

@abhi-mosaic While the llama.cpp guys are working on solving their issue with 16 experts and not 8, I was thinking to quantize 4 bits with the native huggingface BitsAndBytes, still getting an error
P.S. This will enable many people with much smaller computational power still, 66GB to run the model. Single H100 or A100, instead of the current 4.

[abhi-mosaic]

@simsim314 take a look at this comment, I think someone found a workaround by:

• remapping the safetensor weight files to look like separate linear weights, 1 per expert
• edit modeling source code to use a List[nn.Linear]
• new model should work with bitsnbytes

https://huggingface.co/databricks/dbrx-instruct/discussions/10#660566f14f41c0c7c0e54ab9

[peterhgruber]

The model is ~132B params so I think the expected memory usage is:

• 264GB with 16-bit, already possible using our HF quickstart code: https://huggingface.co/databricks/dbrx-instruct#quickstart-guide
• expected 132GB with 8-bit
• expected 99GB with 6-bit
• expected 66GB with 4-bit

Not quite ... they say only 36B parameters are "active on any input", as it is a mixture of experts model.

[wrapss]

Not quite ... they say only 36B parameters are "active on any input", as it is a mixture of experts model.

but the entire model needs to be loaded into memory even if the parameters are not activated

[MohamedAliRashad]

I have the model downloaded in my server if something is added, i can help testing.

[RodriMora]

I have the model downloaded too and can help testing: