Add DBRX Model
Author: abhi-mosaic
Pull request: yes
Labels: name: New model

# What does this PR do?

Add support for `DbrxConfig`, `DbrxModel`, and `DbrxForCausalLM`.

https://huggingface.co/databricks/dbrx-base
https://huggingface.co/databricks/dbrx-instruct

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
Duplicate work here too https://github.com/huggingface/transformers/pull/29910

Hey all! I'll be handling reviews for this model from here - it looks like you have enough to work with for now, but if you get stuck at any point, or you'd like help, please let me know! We're quite excited to get DBRX ported, so don't be shy about reaching out. You can ping me here, or I'm also Matt in the collaboration channel on Slack

Hello, there are some issues in the original implementation of DBRX which I see is still in the codebase of this pull request, the issues only come out during training so I think they're hard to notice currently.
First the function `_gradient_checkpointing_func` called at the `forward` function under `DbrxModel` doesn't support the use of kwargs, so it will throw an error if gradient checkpointing is enabled.
Second, the original repo(And in this pull request), the weights of the MLP experts are in a single tensor, this causes the VRAM usage to explode and go `CUDA Out of Memory` during `backward` through the MLP layer, it can be fixed by separating the weights into different tensors(Into a module list of linear layers). **Unfortunately this fix requires the weights to be converted thus will make it incompatible with the weights in the original huggingface repo of DBRX.**
I tried these 2 fixes on my own implementation [here](https://huggingface.co/v2ray/dbrx-base-fixed), which I tested by fine-tuning it using LoRA with DeepSpeed ZeRO 3, and it worked fine without the above 2 issues.

> it using LoRA with DeepSpeed ZeRO 3

btw, did you test quantized LoRA or keep the base weights in fp16/bf16?

@winglian It's not a quantized LoRA because DeepSpeed ZeRO 3 doesn't support BnB quantization.
Yeah also I didn't test QLoRA but I think it should just work.

All tests pass for now, though I've got to add some more tests from 
```python
from ...generation.test_utils import GenerationTesterMixin
```

Also, I commented out the doc-test for `DbrxForCausalLM.forward()`.

This doc-test is
```python
>>> from transformers import AutoTokenizer, DbrxForCausalLM

>>> model = DbrxForCausalLM.from_pretrained("databricks/dbrx-instruct")
>>> tokenizer = AutoTokenizer.from_pretrained("databricks/dbrx-instruct", trust_remote_code=True)

>>> prompt = "Hey, are you conscious? Can you talk to me?"
>>> inputs = tokenizer(prompt, return_tensors="pt")

>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
```

This doc-test fails because `"databricks/dbrx-instruct"` is a gated-repo.

How can I grant circle.ci access to `"databricks/dbrx-instruct"`? Is there some place in circle.ci where I can pass a token? Is there a standard way of doing this? @Rocketknight1 

@LagPixelLOL Gradient checkpointing should now work. We are working on fixing the fused weights of the MLP experts.

Hi @eitanturok, `dbrx-instruct` is way too big for our doctest runner to load either way! We recommend replacing the `>>>` on each line with `>>` - this looks about the same for users, but our doctest runner won't try to run it. This is the approach we use for other architectures that don't have small checkpoints suitable for the doctest runner - you can see an example in the modeling code for [Command-R](https://github.com/huggingface/transformers/blob/main/src/transformers/models/cohere/modeling_cohere.py#L1074).

Been following work here and in other places to train this model.  

This is the most feature-rich implementation, thanks for that.

However, this implementation(at this time) does not support bitsandbytes int4 quantization.  My understanding for that is that DataBricks did not use nn.linear layers for their model and instead used something else.

I have seen other model uploads that claim to fix this(and they do load with less memory), but those models do not work well with other software in the huggingface ecosystem.

In short, are you aware of the fact the bitsandbytes(and by extension Qlora) is not working with these models weights and is it planned to convert the weights to add support?

> Been following work here and in other places to train this model.
> 
> This is the most feature-rich implementation, thanks for that.
> 
> However, this implementation(at this time) does not support bitsandbytes int4 quantization. My understanding for that is that DataBricks did not use nn.linear layers for their model and instead used something else.
> 
> I have seen other model uploads that claim to fix this(and they do load with less memory), but those models do not work well with other software in the huggingface ecosystem.
> 
> In short, are you aware of the fact the bitsandbytes(and by extension Qlora) is not working with these models weights and is it planned to convert the weights to add support?

Yes, we are working on addressing this.

> > Been following work here and in other places to train this model.
> > This is the most feature-rich implementation, thanks for that.
> > However, this implementation(at this time) does not support bitsandbytes int4 quantization. My understanding for that is that DataBricks did not use nn.linear layers for their model and instead used something else.
> > I have seen other model uploads that claim to fix this(and they do load with less memory), but those models do not work well with other software in the huggingface ecosystem.
> > In short, are you aware of the fact the bitsandbytes(and by extension Qlora) is not working with these models weights and is it planned to convert the weights to add support?
> 
> Yes, we are working on addressing this.

Great!  Really looking forward to it!

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_29921). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Hey @eitanturok, tests are almost there! I pushed one last change to the toctree and now doctests pass. The remaining test failures are: 

- `test_new_cache_format`: We haven't ported LLaMA-like models to this format yet. For now, it's okay to skip this test! We'll update the DBRX code too when we do this update for related models.
- `test_left_padding_compatibility`: This test is failing because `config.pad_token_id` is not set. Once that's set correctly then it should pass!

Thank you to @gante for the diagnosis on those test failures!

@Rocketknight1 Just fixed the errors you mentioned.

Hi, I have a question.

Why are we able to use DBRX with `AutoModelForCausalLM.from_pretrained()` when the model architecture hasn't yet hit the transformers lib? I understand that the weights are available for download, but doesn't the model's `forward()` method need to be in transformers for it to work?

EDIT: This is on transformers 4.39.3

> Hi, I have a question.
> 
> Why are we able to use DBRX with `AutoModelForCausalLM.from_pretrained()` when the model architecture hasn't yet hit the transformers lib? I understand that the weights are available for download, but doesn't the model's `forward()` method need to be in transformers for it to work?
> 
> EDIT: This is on transformers 4.39.3

trust_remote_code=True allows this.  The models out there now have implementations uploaded to their model folder and it uses that code to run the model.

Thank you for that explanation.

If anyone else came to this thread b/c they wanted to check out the model's architecture, and was wondering where modeling_dbrx is, it's here: https://github.com/databricks/dbrx/blob/main/model/modeling_dbrx.py

cc @amyeroberts @ArthurZucker for core maintainer review at this point! (The failing test is flaky and unrelated to this PR)

> Now the biggest blocker is the nn.parameter and the many 2 classes that are wrapping around linear layers + the research code paths.

@ArthurZucker my teammates are prototyping a solution with state dict hooks that should make it possible to (optionally) load the existing weight files into a `DbrxExpertGLUSplit` module with a `List[nn.Linear]`, see PR in progress here: https://github.com/abhi-mosaic/transformers/pull/2/files

I think this would resolve the issue for memory usage for training. Quantization however is still a little bit tricky due to the sequencing of events in `.from_pretrained()`, I think after load, if you manually call quantization on the model with `nn.Linear`, it would work. But we're having trouble making it work out of the box with `device_map=auto` and `quantization_config`, and the state dict hooks.

Btw here is the PR where we are attempting to add the save/load state_dict functionality: https://github.com/abhi-mosaic/transformers/pull/3

Also, one of the community members noted that `past_key_values` is plural in some places but singular in others, this is based on the `modeling_llama.py` code. Is this intentional? Or should it be standardized one way? 

https://huggingface.co/databricks/dbrx-base/discussions/17#660df597d0a9664d8bcdc5f1

> Also, one of the community members noted that `past_key_values` is plural in some places but singular in others, this is based on the `modeling_llama.py` code. Is this intentional? Or should it be standardized one way?
> 
> https://huggingface.co/databricks/dbrx-base/discussions/17#660df597d0a9664d8bcdc5f1

There are some inconsistencies here in `transformers`, but the most common approach is that `past_key_values` is the argument to the top level class like `DbrxModel`, and `past_key_value` is a single element of `past_key_values` that gets passed to a single attention layer. The exact details of what you do inside the model don't matter so much, but I do think `past_key_values` as the top-level kwarg name is needed!

What’s the current status of the PR @abhi-mosaic @Rocketknight1 @cli99? Is it pending the merge of https://github.com/abhi-mosaic/transformers/pull/3? We’ve tried to run that code during training and found some issues. Happy to assist however we can.

Hey @michaelroyzen , we are closing the state dict hook PR (https://github.com/abhi-mosaic/transformers/pull/3) as it would would have added too much complexity. Instead we will be cleaning up this PR, keeping the fused weights in the safetensors files, and adding an option to use `megablocks` (https://github.com/databricks/megablocks) in the DBRX modeling code to improve memory usage during finetuning.

What’s the plan for Megablocks support? The model is quite difficult to fine-tune in its current state even on 32+ H100s due to high memory usage @abhi-databricks @eitanturok 

@michaelroyzen, this HF forum thread helped me fine tune with lower VRAM requirements while waiting for this PR to land: https://huggingface.co/databricks/dbrx-instruct/discussions/10#660566f14f41c0c7c0e54ab9

Specifically, the [SinclairSchneider/dbrx-instruct-quantization-fixed](https://huggingface.co/SinclairSchneider/dbrx-instruct-quantization-fixed) model, posted on that thread, seems to be working well.

@ArthurZucker @BenjaminBossan @winglian  Can we get a final review?

@Rocketknight1 Just addressed all your comments. Would love for you to approve if you feel ready. Cheers!

I'm happy with everything at this point! cc @arthurzucker for core maintainer final review and merging!

One final fix: I just noticed one of the slow tests was loading the entire `dbrx-instruct` model! This is way too big for our test runners. I made a smaller random checkpoint with the same architecture for the test instead and uploaded it at `Rocketknight1/dbrx-tiny-random`.

Hey all! With apologies for the delay - due to the ongoing torrent of model releases, @amyeroberts will be handling the core maintainer review. I've also done a deep review and I'm happy to approve the model at this point, with one final thing I overlooked: The model is missing `@add_start_docstrings` and `@add_start_docstrings_to_model_forward` / `@replace_return_docstrings` on the main model classes, so the documentation will be a bit bare! You can copy those in from another model like `Mixtral`, with any changes you think are appropriate for DBRX.

Hopefully we can get this merged ASAP, though - @amyeroberts let me know if you find any issues other than the docstring decorators!

The failing tests are quite weird and possible connection errors. I'm rerunning on CI which should hopefully resolve 🤞 

Addressed your comments! Let's do the final push and get this in!

> Only thing left is resolving the conflicts with `main` for `docs/source/en/perf_infer_gpu_one.md` and `docs/source/en/tasks/language_modeling.md`.
> 
> As the tests are looking good -- quality checks now pass -- I'd recommend resolving the conflicts before the CI finished. The test suite is unfortunately pretty slow and we can be pretty confident the tests should all pass, as they were two commits ago and the changes have been minimal

Yes, I was just hoping on the train. I'm taking care of this now.

Just merged the conflicts. Now we just let the tests run and hope for the best!

Awesome! Thanks all, merging!

As a follow up, when can we expect to see DBRX updated in the `transformers` documentation?

@eitanturok the documentation is automatically generated from the docstrings / `.md` files in this PR, so you should see it once the release goes out!

Understood. And the release goes out later today, correct?

Yes! The release has just gone out [here](https://github.com/huggingface/transformers/releases/tag/v4.40.0). The frontend docs will be updated as soon as the doc builder finishes, which will happen later today. In the meantime, you can preview the docs on the `main` doc branch [here](https://huggingface.co/docs/transformers/main/en/model_doc/dbrx).