Repository navigation
Add support for DBRX models: dbrx-base and dbrx-instruct #6344
Description
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
Activity
abhi-mosaic commentedon Mar 27, 2024
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:
DbrxExpertGLUlayer are fused along the experts , akatransformer.blocks.12.ffn.experts.mlp.w1, nottransformer.blocks.12.ffn.experts.mlp.list_w1.0.weight. This is what you will find in the.safetensorsweight files, you can see the mapping explicitly here in themodel.safetensors.index.json: https://huggingface.co/databricks/dbrx-instruct/blob/main/model.safetensors.index.json...mlp.w1) after they are loaded from disk.DbrxExpertGLU: https://huggingface.co/databricks/dbrx-instruct/blob/464e701f50aef4c1b59c81fb5667819a5d08e108/modeling_dbrx.py#L749Please let me know if you have any questions!
abhi-mosaic commentedon Mar 27, 2024
The model is ~132B params so I think the expected memory usage is:
ggerganov commentedon Mar 27, 2024
@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 commentedon Mar 27, 2024
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 commentedon Mar 27, 2024
Same, got big servers with very fast and plentyful ram channels, so can try on CPU all the sizes.
sirus20x6 commentedon Mar 28, 2024
Put me in coach, I'm ready to play, today.
veryvanya commentedon Mar 28, 2024
happy to test on my server
simsim314 commentedon Mar 28, 2024
@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 commentedon Mar 28, 2024
@simsim314 take a look at this comment, I think someone found a workaround by:
List[nn.Linear]https://huggingface.co/databricks/dbrx-instruct/discussions/10#660566f14f41c0c7c0e54ab9
peterhgruber commentedon Mar 29, 2024
Not quite ... they say only 36B parameters are "active on any input", as it is a mixture of experts model.
wrapss commentedon Mar 29, 2024
but the entire model needs to be loaded into memory even if the parameters are not activated
MohamedAliRashad commentedon Mar 29, 2024
I have the model downloaded in my server if something is added, i can help testing.
RodriMora commentedon Apr 1, 2024
I have the model downloaded too and can help testing:
nkeilar commentedon Apr 2, 2024
Have a dual 3090 setup and interested in quant to 2bit to see if it will fit in 48GB VRAM, could also test with CPU layers offloaded as running 14900KS. Eric was able to get a The Professor, a 155 Billion parameter model into being able to run on a dual 3090.
22 remaining items
phymbert commentedon Apr 7, 2024
DbrxForCausalLM#6515ggerganov commentedon Apr 8, 2024
No need to upload it - in
ggml-orgwe only want to have models that are used by the CI or for other kind of test/demo purposesphymbert commentedon Apr 8, 2024
Noted, deleted
model: support arch `DbrxForCausalLM` (#6515)
model: support arch `DbrxForCausalLM` (ggml-org#6515)
model: support arch `DbrxForCausalLM` (ggml-org#6515)
model: support arch `DbrxForCausalLM` (ggml-org#6515)
model: support arch `DbrxForCausalLM` (ggml-org#6515)
model: support arch `DbrxForCausalLM` (ggml-org#6515)
model: support arch `DbrxForCausalLM` (ggml-org#6515)