lrl-modelcloud commited on
Commit
f66f594
1 Parent(s): 7cdf1b0

f90ce0b8fdbca174da1b904d6f72ba4c2624cd89bbe3461d64e899f5b7f65f03

Browse files
LICENSE.txt ADDED
@@ -0,0 +1,176 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Databricks Open Model License
2
+
3
+ By using, reproducing, modifying, distributing, performing or displaying
4
+ any portion or element of DBRX or DBRX Derivatives, or otherwise accepting
5
+ the terms of this Agreement, you agree to be bound by this Agreement.
6
+
7
+ Version Release Date: March 27, 2024
8
+
9
+
10
+ Section 1: Definitions
11
+
12
+ “Agreement” means these terms and conditions that govern the use, reproduction,
13
+ modification, distribution, performance or display of DBRX and/or DBRX
14
+ Derivatives and any terms and conditions incorporated by reference.
15
+
16
+ “Databricks” or “we” means Databricks, Inc.
17
+
18
+ “Licensee” or “you” means you, or your employer or any other person or entity
19
+ (if you are entering into this Agreement on such person or entity’s behalf),
20
+ of the age required under applicable laws, rules or regulations to provide
21
+ legal consent and that has legal authority to bind your employer or such other
22
+ person or entity if you are entering in this Agreement on their behalf.
23
+
24
+ “DBRX Derivatives” means all (i) modifications to DBRX, (ii) works based on
25
+ DBRX and (iii) any other derivative works thereof. Outputs are not deemed DBRX
26
+ Derivatives.
27
+
28
+ “DBRX” means the foundational large language models and software and
29
+ algorithms, including machine-learning model code, trained model weights,
30
+ inference-enabling code, training-enabling code, fine-tuning enabling code,
31
+ documentation and other elements of the foregoing identified by Databricks at
32
+ https://github.com/databricks/dbrx, regardless of the source that you obtained
33
+ it from.
34
+
35
+ “Output” means the results of operating DBRX or DBRX Derivatives.
36
+
37
+ As used in this Agreement, “including” means “including without limitation.”
38
+
39
+
40
+ Section 2: License Rights and Conditions on Use and Distribution
41
+
42
+ 2.1 Grant of Rights
43
+
44
+ You are granted a non-exclusive, worldwide, non-transferable and royalty-free
45
+ limited license under Databricks’ intellectual property or other rights owned
46
+ by Databricks embodied in DBRX to use, reproduce, distribute, copy, modify,
47
+ and create derivative works of DBRX in accordance with the terms of this
48
+ Agreement.
49
+
50
+ 2.2 Reproduction and Distribution
51
+
52
+ 1. All distributions of DBRX or DBRX Derivatives must be accompanied by a
53
+ "Notice" text file that contains the following notice: "DBRX is provided
54
+ under and subject to the Databricks Open Model License, Copyright ©
55
+ Databricks, Inc. All rights reserved."
56
+
57
+ 2. If you distribute or make DBRX or DBRX Derivatives available to a third
58
+ party, you must provide a copy of this Agreement to such third party.
59
+
60
+ 3. You must cause any modified files that you distribute to carry prominent
61
+ notices stating that you modified the files.
62
+
63
+ You may add your own intellectual property statement to your modifications of
64
+ DBRX and, except as set forth in this Section, may provide additional or
65
+ different terms and conditions for use, reproduction, or distribution of DBRX
66
+ or DBRX Derivatives as a whole, provided your use, reproduction, modification,
67
+ distribution, performance, and display of DBRX or DBRX Derivatives otherwise
68
+ complies with the terms and conditions of this Agreement. Any additional or
69
+ different terms and conditions you impose must not conflict with the terms of
70
+ this Agreement and in the event of a conflict, the terms and conditions of this
71
+ Agreement shall govern over any such additional or different terms and conditions.
72
+
73
+ 2.3 Use Restrictions
74
+
75
+ You will not use DBRX or DBRX Derivatives or any Output to improve any other
76
+ large language model (excluding DBRX or DBRX Derivatives).
77
+
78
+ You will not use DBRX or DBRX Derivatives:
79
+
80
+ 1. for any restricted use set forth in the Databricks Open Model Acceptable
81
+ Use Policy identified at
82
+ https://www.databricks.com/legal/acceptable-use-policy-open-model
83
+ ("Acceptable Use Policy"), which is hereby incorporated by reference into
84
+ this Agreement; or
85
+
86
+ 2. in violation of applicable laws and regulations.
87
+
88
+ To the maximum extent permitted by law, Databricks reserves the right to
89
+ restrict (remotely or otherwise) usage of DBRX or DBRX Derivatives that
90
+ Databricks reasonably believes are in violation of this Agreement.
91
+
92
+
93
+ Section 3: Additional Commercial Terms
94
+
95
+ If, on the DBRX version release date, the monthly active users of the products
96
+ or services made available by or for Licensee, or Licensee’s affiliates, is
97
+ greater than 700 million monthly active users in the preceding calendar month,
98
+ you must request a license from Databricks, which we may grant to you in our
99
+ sole discretion, and you are not authorized to exercise any of the rights under
100
+ this Agreement unless or until Databricks otherwise expressly grants you such
101
+ rights.
102
+
103
+ If you receive DBRX or DBRX Derivatives from a direct or indirect licensee as
104
+ part of an integrated end user product, then this section (Section 3) of the
105
+ Agreement will not apply to you.
106
+
107
+
108
+ Section 4: Additional Provisions
109
+
110
+ 4.1 Updates
111
+
112
+ Databricks may update DBRX from time to time, and you must make reasonable
113
+ efforts to use the latest version of DBRX.
114
+
115
+ 4.2 Intellectual Property
116
+
117
+ a. No trademark licenses are granted under this Agreement, and in connection
118
+ with DBRX or DBRX Derivatives, neither Databricks nor Licensee may use any name
119
+ or mark owned by or associated with the other or any of its affiliates, except
120
+ as required for reasonable and customary use in describing and redistributing
121
+ DBRX or DBRX Derivatives.
122
+
123
+ b. Subject to Databricks’ ownership of DBRX and DRBX Derivatives made by or for
124
+ Databricks, with respect to any DBRX Derivatives that are made by you, as
125
+ between you and Databricks, you are and will be the owner of such DBRX
126
+ Derivatives.
127
+
128
+ c. Databricks claims no ownership rights in Outputs. You are responsible for
129
+ Outputs and their subsequent uses.
130
+
131
+ d. If you institute litigation or other proceedings against Databricks or any
132
+ entity (including a cross-claim or counterclaim in a lawsuit) alleging that
133
+ DBRX or Outputs or results therefrom, or any portion of any of the foregoing,
134
+ constitutes infringement of intellectual property or other rights owned or
135
+ licensable by you, then any licenses granted to you under this Agreement shall
136
+ terminate as of the date such litigation or claim is filed or instituted. You
137
+ will indemnify and hold harmless Databricks from and against any claim by any
138
+ third party arising out of or related to your use or distribution of DBRX or
139
+ DBRX Derivatives.
140
+
141
+ 4.3 DISCLAIMER OF WARRANTY
142
+
143
+ UNLESS REQUIRED BY APPLICABLE LAW, DBRX AND ANY OUTPUT AND RESULTS THEREFROM
144
+ ARE PROVIDED ON AN “AS IS” BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER
145
+ EXPRESS OR IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OF TITLE,
146
+ NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU
147
+ ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING OR
148
+ REDISTRIBUTING DBRX OR DBRX DERIVATIVES AND ANY OUTPUT AND ASSUME ANY RISKS
149
+ ASSOCIATED WITH YOUR USE OF DBRX OR DBRX DERIVATIVES AND ANY OUTPUT AND RESULTS.
150
+
151
+ 4.4 LIMITATION OF LIABILITY
152
+
153
+ IN NO EVENT WILL DATABRICKS OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF
154
+ LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR
155
+ OTHERWISE, ARISING OUT OF THIS AGREEMENT, FOR ANY LOST PROFITS OR ANY INDIRECT,
156
+ SPECIAL, CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN IF
157
+ DATABRICKS OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE
158
+ FOREGOING.
159
+
160
+ 4.5 Term and Termination
161
+
162
+ The term of this Agreement will commence upon your acceptance of this Agreement
163
+ or access to DBRX or DBRX Derivatives and will continue in full force and
164
+ effect until terminated in accordance with the terms and conditions herein.
165
+ Databricks may terminate this Agreement if you are in breach of any term or
166
+ condition of this Agreement. Upon termination of this Agreement, you shall
167
+ delete and cease use of DBRX or any DBRX Derivatives. Sections 1, 4.2(d), 4.3,
168
+ 4.4, and 4.6 shall survive the termination of this Agreement.
169
+
170
+ 4.6 Governing Law and Jurisdiction
171
+
172
+ This Agreement will be governed and construed under the laws of the State of
173
+ California without regard to choice of law principles, and the UN Convention
174
+ on Contracts for the International Sale of Goods does not apply to this
175
+ Agreement. The courts of California shall have exclusive jurisdiction of any
176
+ dispute arising out of this Agreement.
NOTICE.txt ADDED
@@ -0,0 +1 @@
 
 
1
+ DBRX is provided under and subject to the Databricks Open Model License, Copyright © Databricks, Inc. All rights reserved.
README.md CHANGED
@@ -1,3 +1,172 @@
1
- ---
2
- license: unknown
3
- ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ extra_gated_heading: You need to share contact information with Databricks to access this model
3
+ extra_gated_prompt: >-
4
+
5
+ ### DBRX Terms of Use
6
+
7
+ Use of DBRX is governed by the [Databricks Open Model License](https://www.databricks.com/legal/open-model-license) and the [Databricks Open Model Acceptable Use Policy](https://www.databricks.com/legal/acceptable-use-policy-open-model).
8
+
9
+ extra_gated_fields:
10
+ First Name: text
11
+ Last Name: text
12
+ Organization: text
13
+ Purpose for Base Model Access: text
14
+ By clicking 'Submit' below, I accept the terms of the license and acknowledge that the information I provide will be collected, stored, processed, and shared in accordance with Databricks' Privacy Notice and I understand I can update my preferences at any time: checkbox
15
+ extra_gated_description: >-
16
+ The information you provide will be collected, stored, processed, and shared in accordance with Databricks [Privacy Notice](https://www.databricks.com/legal/privacynotice).
17
+ extra_gated_button_content: Submit
18
+ inference: false
19
+ license: other
20
+ license_name: databricks-open-model-license
21
+ license_link: https://www.databricks.com/legal/open-model-license
22
+ ---
23
+
24
+ # DBRX Base
25
+
26
+ * DBRX Base is a mixture-of-experts (MoE) large language model trained from scratch by Databricks.
27
+ * We are releasing both DBRX Base, a pretrained base model, and DBRX Instruct, a fine-tuned version for few-turn interactions, under [an open license](https://www.databricks.com/legal/open-model-license).
28
+ * This is the repository for DBRX Base. DBRX Instruct can be found [here](https://huggingface.co/databricks/dbrx-instruct).
29
+ * For full details on the DBRX models, please read our [technical blog post](https://www.databricks.com/blog/introducing-dbrx-new-state-art-open-llm).
30
+
31
+
32
+ ## Model Overview
33
+ DBRX is a [transformer-based](https://www.isattentionallyouneed.com/) decoder-only large language model (LLM) that was trained using next-token prediction.
34
+ It uses a *fine-grained* mixture-of-experts (MoE) architecture with 132B total parameters of which 36B parameters are active on any input.
35
+ It was pre-trained on 12T tokens of text and code data.
36
+ Compared to other open MoE models like Mixtral-8x7B and Grok-1, DBRX is fine-grained, meaning it uses a larger number of smaller experts. DBRX has 16 experts and chooses 4, while Mixtral-8x7B and Grok-1 have 8 experts and choose 2.
37
+ This provides 65x more possible combinations of experts and we found that this improves model quality.
38
+ DBRX uses rotary position encodings (RoPE), gated linear units (GLU), and grouped query attention (GQA).
39
+ It uses the GPT-4 tokenizer as provided in the [tiktoken](https://github.com/openai/tiktoken) repository.
40
+ We made these choices based on exhaustive evaluation and scaling experiments.
41
+
42
+ DBRX was pretrained on 12T tokens of carefully curated data and a maximum context length of 32K tokens.
43
+ We estimate that this data is at least 2x better token-for-token than the data we used to pretrain the MPT family of models.
44
+ This new dataset was developed using the full suite of Databricks tools, including Apache Spark™ and Databricks notebooks for data processing, and Unity Catalog for data management and governance.
45
+ We used curriculum learning for pretraining, changing the data mix during training in ways we found to substantially improve model quality.
46
+
47
+ * **Inputs:** DBRX only accepts text-based inputs and accepts a context length of up to 32768 tokens.
48
+ * **Outputs:** DBRX only produces text-based outputs.
49
+ * **Model Architecture:** More detailed information about DBRX Instruct and DBRX Base can be found in our [technical blog post](https://www.databricks.com/blog/introducing-dbrx-new-state-art-open-llm).
50
+ * **License:** [Databricks Open Model License](https://www.databricks.com/legal/open-model-license)
51
+ * **Acceptable Use Policy:** [Databricks Open Model Acceptable Use Policy](https://www.databricks.com/legal/acceptable-use-policy-open-model)
52
+ * **Version:** 1.0
53
+ * **Owner:** Databricks, Inc.
54
+
55
+
56
+ ## Usage
57
+ These are several general ways to use the DBRX models:
58
+ * DBRX Base and DBRX Instruct are available for download on HuggingFace (see our Quickstart guide below). This is the HF repository for DBRX Base; DBRX Instruct can be found [here](https://huggingface.co/databricks/dbrx-instruct).
59
+ * The DBRX model repository can be found on GitHub [here](https://github.com/databricks/dbrx).
60
+ * DBRX Base and DBRX Instruct are available with [Databricks Foundation Model APIs](https://docs.databricks.com/en/machine-learning/foundation-models/index.html) via both *Pay-per-token* and *Provisioned Throughput* endpoints. These are enterprise-ready deployments.
61
+ * For more information on how to fine-tune using LLM-Foundry, please take a look at our LLM pretraining and fine-tuning [documentation](https://github.com/mosaicml/llm-foundry/blob/main/scripts/train/README.md).
62
+
63
+
64
+ ## Quickstart Guide
65
+ **NOTE: This is DBRX Base, and has not been instruction finetuned. It has not been trained for interactive chat and is only a completion model.**
66
+ If you are looking for the finetuned model, please use [DBRX Instruct](https://huggingface.co/databricks/dbrx-instruct).
67
+
68
+ Getting started with DBRX models is easy with the `transformers` library. The model requires ~264GB of RAM and the following packages:
69
+
70
+ ```bash
71
+ pip install "transformers>=4.39.2" "tiktoken>=0.6.0"
72
+ ```
73
+
74
+ If you'd like to speed up download time, you can use the `hf_transfer` package as described by Huggingface [here](https://huggingface.co/docs/huggingface_hub/en/guides/download#faster-downloads).
75
+ ```bash
76
+ pip install hf_transfer
77
+ export HF_HUB_ENABLE_HF_TRANSFER=1
78
+ ```
79
+
80
+ You will need to request access to this repository to download the model. Once this is granted,
81
+ [obtain an access token](https://huggingface.co/docs/hub/en/security-tokens) with `read` permission, and supply the token below.
82
+
83
+ ### Run the model on a CPU:
84
+ ```python
85
+ from transformers import AutoTokenizer, AutoModelForCausalLM
86
+ import torch
87
+
88
+ tokenizer = AutoTokenizer.from_pretrained("databricks/dbrx-base", trust_remote_code=True, token="hf_YOUR_TOKEN")
89
+ model = AutoModelForCausalLM.from_pretrained("databricks/dbrx-base", device_map="cpu", torch_dtype=torch.bfloat16, trust_remote_code=True, token="hf_YOUR_TOKEN")
90
+
91
+ input_text = "Databricks was founded in "
92
+ input_ids = tokenizer(input_text, return_tensors="pt")
93
+
94
+ outputs = model.generate(**input_ids, max_new_tokens=100)
95
+ print(tokenizer.decode(outputs[0]))
96
+ ```
97
+
98
+ ### Run the model on multiple GPUs:
99
+ ```python
100
+ from transformers import AutoTokenizer, AutoModelForCausalLM
101
+ import torch
102
+
103
+ tokenizer = AutoTokenizer.from_pretrained("databricks/dbrx-base", trust_remote_code=True, token="hf_YOUR_TOKEN")
104
+ model = AutoModelForCausalLM.from_pretrained("databricks/dbrx-base", device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True, token="hf_YOUR_TOKEN")
105
+
106
+ input_text = "Databricks was founded in "
107
+ input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
108
+
109
+ outputs = model.generate(**input_ids, max_new_tokens=100)
110
+ print(tokenizer.decode(outputs[0]))
111
+ ```
112
+ If your GPU system supports [FlashAttention2](https://huggingface.co/docs/transformers/perf_infer_gpu_one#flashattention-2), you can add `attn_implementation=”flash_attention_2”` as a keyword to `AutoModelForCausalLM.from_pretrained()` to achieve faster inference.
113
+
114
+
115
+ ## Limitations and Ethical Considerations
116
+ ### Training Dataset Limitations
117
+ The DBRX models were trained on 12T tokens of text, with a knowledge cutoff date of December 2023.
118
+
119
+ The training mix used for DBRX contains both natural-language and code examples. The vast majority of our training data is in the English language. We did not test DBRX for non-English proficiency. Therefore, DBRX should be considered a generalist model for text-based use in the English language.
120
+
121
+ DBRX does not have multimodal capabilities.
122
+
123
+ ### Associated Risks and Recommendations
124
+ All foundation models are novel technologies that carry various risks, and may output information that is inaccurate, incomplete, biased, or offensive.
125
+ Users should exercise judgment and evaluate such output for accuracy and appropriateness for their desired use case before using or sharing it.
126
+ Databricks recommends [using retrieval augmented generation (RAG)](https://www.databricks.com/glossary/retrieval-augmented-generation-rag) in scenarios where accuracy and fidelity are important.
127
+ We also recommend that anyone using or fine-tuning either DBRX Base or DBRX Instruct perform additional testing around safety in the context of their particular application and domain.
128
+
129
+
130
+ ## Intended Uses
131
+ ### Intended Use Cases
132
+ The DBRX models are open, general-purpose LLMs intended and licensed for both commercial and research applications.
133
+ They can be further fine-tuned for various domain-specific natural language and coding tasks.
134
+ DBRX Base can be used as an off-the-shelf model for text completion for general English-language and coding tasks.
135
+
136
+ Please review the Associated Risks section above, as well as the [Databricks Open Model License](https://www.databricks.com/legal/open-model-license) and [Databricks Open Model Acceptable Use Policy](https://www.databricks.com/legal/acceptable-use-policy-open-model) for further information about permissible uses of DBRX Base and its derivatives.
137
+
138
+ ### Out-of-Scope Use Cases
139
+ DBRX models are not intended to be used out-of-the-box in non-English languages and do not support native code execution, or other forms of function-calling.
140
+ DBRX models should not be used in any manner that violates applicable laws or regulations or in any other way that is prohibited by the [Databricks Open Model License](https://www.databricks.com/legal/open-model-license) and [Databricks Open Model Acceptable Use Policy](https://www.databricks.com/legal/acceptable-use-policy-open-model).
141
+
142
+
143
+ ## Training Stack
144
+ MoE models are complicated to train, and the training of DBRX Base and DBRX Instruct was heavily supported by Databricks’ infrastructure for data processing and large-scale LLM training (e.g., [Composer](https://github.com/mosaicml/composer), [Streaming](https://github.com/mosaicml/streaming), [Megablocks](https://github.com/stanford-futuredata/megablocks), and [LLM Foundry](https://github.com/mosaicml/llm-foundry)).
145
+
146
+ Composer is our core library for large-scale training.
147
+ It provides an optimized training loop, easy [checkpointing](https://docs.mosaicml.com/projects/composer/en/latest/trainer/checkpointing.html) and [logging](https://docs.mosaicml.com/projects/composer/en/latest/trainer/logging.html#wood-logging),
148
+ [FSDP](https://pytorch.org/docs/stable/fsdp.html)-based [model sharding](https://docs.mosaicml.com/projects/composer/en/latest/notes/distributed_training.html#fullyshardeddataparallel-fsdp),
149
+ convenient [abstractions](https://docs.mosaicml.com/projects/composer/en/latest/trainer/time.html), extreme customizability via [callbacks](https://docs.mosaicml.com/projects/composer/en/latest/trainer/callbacks.html), and more.
150
+
151
+ Streaming enables fast, low cost, and scalable training on large datasets from cloud storage. It handles a variety of challenges around deterministic resumption as node counts change, avoiding redundant downloads across devices, high-quality shuffling at scale, sample-level random access, and speed.
152
+
153
+ Megablocks is a lightweight library for MoE training. Crucially, it supports “dropless MoE,” which avoids inefficient padding and is intended to provide deterministic outputs for a given sequence no matter what other sequences are in the batch.
154
+
155
+ LLM Foundry ties all of these libraries together to create a simple LLM pretraining, fine-tuning, and inference experience.
156
+
157
+ DBRX was trained using proprietary optimized versions of the above open source libraries, along with our [LLM training platform](https://www.databricks.com/product/machine-learning/mosaic-ai-training).
158
+
159
+
160
+ ## Evaluation
161
+ We find that DBRX outperforms established open-source and open-weight base models on the [Databricks Model Gauntlet](https://www.databricks.com/blog/llm-evaluation-for-icl), the [Hugging Face Open LLM Leaderboard](https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard), and HumanEval.
162
+ The Databricks Model Gauntlet measures performance on more than 30 tasks across six categories: world knowledge, common sense reasoning, language understanding, reading comprehension, symbolic problem solving, and programming.
163
+ The Hugging Face Open LLM Leaderboard measures the average of ARC-Challenge, HellaSwag, MMLU, TruthfulQA, Winogrande and GSM8k.
164
+ HumanEval measures coding ability.
165
+
166
+ Full evaluation details can be found in our [technical blog post](https://www.databricks.com/blog/introducing-dbrx-new-state-art-open-llm).
167
+
168
+
169
+ ## Acknowledgements
170
+ The DBRX models were made possible thanks in large part to the open-source community, especially:
171
+ * The [MegaBlocks](https://arxiv.org/abs/2211.15841) library, which established a foundation for our MoE implementation.
172
+ * [PyTorch FSDP](https://arxiv.org/abs/2304.11277), which we built on for distributed training.
config.json ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "DbrxForCausalLM"
4
+ ],
5
+ "attn_config": {
6
+ "clip_qkv": 8,
7
+ "kv_n_heads": 8,
8
+ "model_type": "",
9
+ "rope_theta": 500000
10
+ },
11
+ "auto_map": {
12
+ "AutoConfig": "configuration_dbrx.DbrxConfig",
13
+ "AutoModelForCausalLM": "modeling_dbrx.DbrxForCausalLM"
14
+ },
15
+ "d_model": 6144,
16
+ "emb_pdrop": 0.0,
17
+ "ffn_config": {
18
+ "ffn_hidden_size": 10752,
19
+ "model_type": "",
20
+ "moe_jitter_eps": 0.01,
21
+ "moe_loss_weight": 0.05,
22
+ "moe_num_experts": 16,
23
+ "moe_top_k": 4
24
+ },
25
+ "initializer_range": 0.02,
26
+ "max_seq_len": 32768,
27
+ "model_type": "dbrx_converted",
28
+ "n_heads": 48,
29
+ "n_layers": 40,
30
+ "output_router_logits": false,
31
+ "resid_pdrop": 0.0,
32
+ "router_aux_loss_coef": 0.05,
33
+ "tie_word_embeddings": false,
34
+ "torch_dtype": "bfloat16",
35
+ "transformers_version": "4.38.2",
36
+ "use_cache": true,
37
+ "vocab_size": 100352
38
+ }
configuration_dbrx.py ADDED
@@ -0,0 +1,264 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Dbrx configuration."""
2
+ from typing import Any, Optional
3
+
4
+ from transformers.configuration_utils import PretrainedConfig
5
+ from transformers.utils import logging
6
+
7
+ logger = logging.get_logger(__name__)
8
+
9
+ DBRX_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
10
+
11
+
12
+ class DbrxAttentionConfig(PretrainedConfig):
13
+ """Configuration class for Dbrx Attention.
14
+
15
+ [`DbrxAttention`] class. It is used to instantiate attention layers
16
+ according to the specified arguments, defining the layers architecture.
17
+
18
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
19
+ documentation from [`PretrainedConfig`] for more information.
20
+
21
+ Args:
22
+ attn_pdrop (`float`, *optional*, defaults to 0.0):
23
+ The dropout probability for the attention layers.
24
+ clip_qkv (`float`, *optional*, defualts to None):
25
+ If not `None`, clip the queries, keys, and values in the attention layer to this value.
26
+ kv_n_heads (Optional[int]): For grouped_query_attention only, allow user to specify number of kv heads.
27
+ rope_theta (float): The base frequency for rope.
28
+ """
29
+
30
+ def __init__(
31
+ self,
32
+ attn_pdrop: float = 0,
33
+ clip_qkv: Optional[float] = None,
34
+ kv_n_heads: int = 1,
35
+ rope_theta: float = 10000.0,
36
+ **kwargs: Any,
37
+ ):
38
+ super().__init__(**kwargs)
39
+ self.attn_pdrop = attn_pdrop
40
+ self.clip_qkv = clip_qkv
41
+ self.kv_n_heads = kv_n_heads
42
+ self.rope_theta = rope_theta
43
+
44
+ for k in ['model_type']:
45
+ if k in kwargs:
46
+ kwargs.pop(k)
47
+ if len(kwargs) != 0:
48
+ raise ValueError(f'Found unknown {kwargs=}')
49
+
50
+ @classmethod
51
+ def from_pretrained(cls, pretrained_model_name_or_path: str,
52
+ **kwargs: Any) -> 'PretrainedConfig':
53
+ cls._set_token_in_kwargs(kwargs)
54
+
55
+ config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path,
56
+ **kwargs)
57
+
58
+ if config_dict.get('model_type') == 'dbrx':
59
+ config_dict = config_dict['attn_config']
60
+
61
+ if 'model_type' in config_dict and hasattr(
62
+ cls,
63
+ 'model_type') and config_dict['model_type'] != cls.model_type:
64
+ logger.warning(
65
+ f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
66
+ +
67
+ f'{cls.model_type}. This is not supported for all configurations of models and can yield errors.'
68
+ )
69
+
70
+ return cls.from_dict(config_dict, **kwargs)
71
+
72
+
73
+ class DbrxFFNConfig(PretrainedConfig):
74
+ """Configuration class for Dbrx FFN.
75
+
76
+ [`DbrxFFN`] class. It is used to instantiate feedforward layers according to
77
+ the specified arguments, defining the layers architecture.
78
+
79
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
80
+ documentation from [`PretrainedConfig`] for more information.
81
+
82
+ Args:
83
+ ffn_act_fn (dict, optional): A dict specifying activation function for the FFN.
84
+ The dict should have a key 'name' with the value being the name of
85
+ the activation function along with any additional keyword arguments.
86
+ ffn_hidden_size (int, optional): The hidden size of the feedforward network.
87
+ moe_num_experts (int, optional): The number of experts in the mixture of experts layer.
88
+ moe_top_k (int, optional): The number of experts to use in the mixture of experts layer.
89
+ moe_jitter_eps (float, optional): The jitter epsilon for the mixture of experts layer.
90
+ moe_loss_weight (float, optional): The loss weight for the mixture of experts layer.
91
+ moe_normalize_expert_weights (float, optional): The normalization factor for the expert weights.
92
+ uniform_expert_assignment (bool, optional): Whether to use uniform expert assignment.
93
+ This should only be used for benchmarking purposes.
94
+ """
95
+
96
+ def __init__(
97
+ self,
98
+ ffn_act_fn: Optional[dict] = None,
99
+ ffn_hidden_size: int = 3584,
100
+ moe_num_experts: int = 4,
101
+ moe_top_k: int = 1,
102
+ moe_jitter_eps: Optional[float] = None,
103
+ moe_loss_weight: float = 0.01,
104
+ moe_normalize_expert_weights: Optional[float] = 1,
105
+ uniform_expert_assignment: bool = False,
106
+ **kwargs: Any,
107
+ ):
108
+ super().__init__()
109
+ if ffn_act_fn is None:
110
+ ffn_act_fn = {'name': 'silu'}
111
+ self.ffn_act_fn = ffn_act_fn
112
+ self.ffn_hidden_size = ffn_hidden_size
113
+ self.moe_num_experts = moe_num_experts
114
+ self.moe_top_k = moe_top_k
115
+ self.moe_jitter_eps = moe_jitter_eps
116
+ self.moe_loss_weight = moe_loss_weight
117
+ self.moe_normalize_expert_weights = moe_normalize_expert_weights
118
+ self.uniform_expert_assignment = uniform_expert_assignment
119
+
120
+ for k in ['model_type']:
121
+ if k in kwargs:
122
+ kwargs.pop(k)
123
+ if len(kwargs) != 0:
124
+ raise ValueError(f'Found unknown {kwargs=}')
125
+
126
+ @classmethod
127
+ def from_pretrained(cls, pretrained_model_name_or_path: str,
128
+ **kwargs: Any) -> 'PretrainedConfig':
129
+ cls._set_token_in_kwargs(kwargs)
130
+
131
+ config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path,
132
+ **kwargs)
133
+
134
+ if config_dict.get('model_type') == 'dbrx':
135
+ config_dict = config_dict['ffn_config']
136
+
137
+ if 'model_type' in config_dict and hasattr(
138
+ cls,
139
+ 'model_type') and config_dict['model_type'] != cls.model_type:
140
+ logger.warning(
141
+ f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
142
+ +
143
+ f'{cls.model_type}. This is not supported for all configurations of models and can yield errors.'
144
+ )
145
+
146
+ return cls.from_dict(config_dict, **kwargs)
147
+
148
+
149
+ class DbrxConfig(PretrainedConfig):
150
+ """Configuration class for Dbrx.
151
+
152
+ [`DbrxModel`]. It is used to instantiate a Dbrx model according to the
153
+ specified arguments, defining the model architecture.
154
+
155
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
156
+ documentation from [`PretrainedConfig`] for more information.
157
+
158
+
159
+ Args:
160
+ d_model (`int`, *optional*, defaults to 6144):
161
+ Dimensionality of the embeddings and hidden states.
162
+ n_heads (`int`, *optional*, defaults to 48):
163
+ Number of attention heads for each attention layer in the Transformer encoder.
164
+ n_layers (`int`, *optional*, defaults to 40):
165
+ Number of hidden layers in the Transformer encoder.
166
+ max_seq_len (`int`, *optional*, defaults to 32768):
167
+ The maximum sequence length of the model.
168
+ vocab_size (`int`, *optional*, defaults to 100352):
169
+ Vocabulary size of the Dbrx model. Defines the maximum number of different tokens that can be represented by
170
+ the `inputs_ids` passed when calling [`DbrxModel`].
171
+ resid_pdrop (`float`, *optional*, defaults to 0.0):
172
+ The dropout probability applied to the attention output before combining with residual.
173
+ emb_pdrop (`float`, *optional*, defaults to 0.0):
174
+ The dropout probability for the embedding layer.
175
+ attn_config (`dict`, *optional*):
176
+ A dictionary used to configure the model's attention module.
177
+ ffn_config (`dict`, *optional*):
178
+ A dictionary used to configure the model's FFN module.
179
+ use_cache (`bool`, *optional*, defaults to `False`):
180
+ Whether or not the model should return the last key/values attentions (not used by all models).
181
+ initializer_range (`float`, *optional*, defaults to 0.02):
182
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
183
+ output_router_logits (`bool`, *optional*, defaults to `False`):
184
+ Whether or not the router logits should be returned by the model. Enabling this will also
185
+ allow the model to output the auxiliary loss. See [here]() for more details
186
+ router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
187
+ The aux loss factor for the total loss.
188
+
189
+
190
+ Example:
191
+ ```python
192
+ >>> from transformers import DbrxConfig, DbrxModel
193
+
194
+ >>> # Initializing a Dbrx configuration
195
+ >>> configuration = DbrxConfig()
196
+
197
+ >>> # Initializing a model (with random weights) from the configuration
198
+ >>> model = DbrxModel(configuration)
199
+
200
+ >>> # Accessing the model configuration
201
+ >>> configuration = model.config
202
+ ```
203
+ """
204
+
205
+ model_type = 'dbrx'
206
+ attribute_map = {
207
+ 'num_attention_heads': 'n_heads',
208
+ 'hidden_size': 'd_model',
209
+ 'num_hidden_layers': 'n_layers',
210
+ 'max_position_embeddings': 'max_seq_len'
211
+ }
212
+
213
+ def __init__(
214
+ self,
215
+ d_model: int = 2048,
216
+ n_heads: int = 16,
217
+ n_layers: int = 24,
218
+ max_seq_len: int = 2048,
219
+ vocab_size: int = 32000,
220
+ resid_pdrop: float = 0.0,
221
+ emb_pdrop: float = 0.0,
222
+ attn_config: Optional[DbrxAttentionConfig] = None,
223
+ ffn_config: Optional[DbrxFFNConfig] = None,
224
+ use_cache: bool = True,
225
+ initializer_range: float = 0.02,
226
+ output_router_logits: bool = False,
227
+ router_aux_loss_coef: float = 0.05,
228
+ **kwargs: Any,
229
+ ):
230
+ if attn_config is None:
231
+ self.attn_config = DbrxAttentionConfig()
232
+ elif isinstance(attn_config, dict):
233
+ self.attn_config = DbrxAttentionConfig(**attn_config)
234
+ else:
235
+ self.attn_config = attn_config
236
+
237
+ if ffn_config is None:
238
+ self.ffn_config = DbrxFFNConfig()
239
+ elif isinstance(ffn_config, dict):
240
+ self.ffn_config = DbrxFFNConfig(**ffn_config)
241
+ else:
242
+ self.ffn_config = ffn_config
243
+
244
+ self.d_model = d_model
245
+ self.n_heads = n_heads
246
+ self.n_layers = n_layers
247
+ self.max_seq_len = max_seq_len
248
+ self.vocab_size = vocab_size
249
+ self.resid_pdrop = resid_pdrop
250
+ self.emb_pdrop = emb_pdrop
251
+ self.use_cache = use_cache
252
+ self.initializer_range = initializer_range
253
+ self.output_router_logits = output_router_logits
254
+ self.router_aux_loss_coef = router_aux_loss_coef
255
+
256
+ tie_word_embeddings = kwargs.pop('tie_word_embeddings', False)
257
+ if tie_word_embeddings:
258
+ raise ValueError(
259
+ 'tie_word_embeddings is not supported for Dbrx models.')
260
+
261
+ super().__init__(
262
+ tie_word_embeddings=tie_word_embeddings,
263
+ **kwargs,
264
+ )
generation_config.json ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "eos_token_id": [
4
+ 100257
5
+ ],
6
+ "transformers_version": "4.38.2"
7
+ }
model.safetensors.index.json ADDED
The diff for this file is too large to render. See raw diff
 
modeling_dbrx.py ADDED
@@ -0,0 +1,1479 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """PyTorch Dbrx model."""
2
+
3
+ import math
4
+ import warnings
5
+ from copy import deepcopy
6
+ from functools import partial
7
+ from typing import Any, Callable, Dict, Optional, Tuple, Union
8
+
9
+ import torch
10
+ import torch.nn.functional as F
11
+ import torch.utils.checkpoint
12
+ from torch import nn
13
+ from transformers.cache_utils import Cache, DynamicCache, StaticCache
14
+ from transformers.modeling_attn_mask_utils import AttentionMaskConverter
15
+ from transformers.modeling_outputs import (MoeCausalLMOutputWithPast,
16
+ MoeModelOutputWithPast)
17
+ from transformers.modeling_utils import PreTrainedModel
18
+ from transformers.utils import is_flash_attn_2_available, logging
19
+
20
+ from .configuration_dbrx import DbrxAttentionConfig, DbrxConfig, DbrxFFNConfig
21
+
22
+ if is_flash_attn_2_available():
23
+ try:
24
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
25
+ from flash_attn.bert_padding import pad_input # noqa
26
+ from flash_attn.bert_padding import index_first_axis, unpad_input
27
+ except:
28
+ pass
29
+
30
+ logger = logging.get_logger(__name__)
31
+
32
+ _CONFIG_FOR_DOC = 'DbrxConfig'
33
+
34
+ #############################################################################
35
+ # Copied from LLaMaRotaryEmbedding
36
+ #############################################################################
37
+
38
+
39
+ class DbrxRotaryEmbedding(nn.Module):
40
+
41
+ def __init__(self,
42
+ dim: int,
43
+ max_position_embeddings: int = 2048,
44
+ base: float = 10000.0,
45
+ scaling_factor: float = 1.0):
46
+ super().__init__()
47
+ self.scaling_factor = scaling_factor
48
+ self.dim = dim
49
+ self.max_position_embeddings = max_position_embeddings
50
+ self.base = base
51
+ inv_freq = 1.0 / (self.base**(
52
+ torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim))
53
+ self.register_buffer('inv_freq', inv_freq, persistent=False)
54
+ # For BC we register cos and sin cached
55
+ self.max_seq_len_cached = max_position_embeddings
56
+
57
+ @torch.no_grad()
58
+ def forward(
59
+ self, x: torch.Tensor, position_ids: torch.LongTensor
60
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
61
+ # x: [bs, num_attention_heads, seq_len, head_size]
62
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(
63
+ position_ids.shape[0], -1, 1)
64
+ position_ids_expanded = position_ids[:, None, :].float()
65
+ # Force float32 since bfloat16 loses precision on long contexts
66
+ # See https://github.com/huggingface/transformers/pull/29285
67
+ device_type = x.device.type
68
+ device_type = device_type if isinstance(
69
+ device_type, str) and device_type != 'mps' else 'cpu'
70
+ with torch.autocast(device_type=device_type, enabled=False):
71
+ freqs = (inv_freq_expanded.float()
72
+ @ position_ids_expanded.float()).transpose(1, 2)
73
+ emb = torch.cat((freqs, freqs), dim=-1)
74
+ cos = emb.cos()
75
+ sin = emb.sin()
76
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
77
+
78
+
79
+ def rotate_half(x: torch.Tensor) -> torch.Tensor:
80
+ """Rotates half the hidden dims of the input."""
81
+ x1 = x[..., :x.shape[-1] // 2]
82
+ x2 = x[..., x.shape[-1] // 2:]
83
+ return torch.cat((-x2, x1), dim=-1)
84
+
85
+
86
+ def apply_rotary_pos_emb(
87
+ q: torch.Tensor,
88
+ k: torch.Tensor,
89
+ cos: torch.Tensor,
90
+ sin: torch.Tensor,
91
+ unsqueeze_dim: int = 1) -> Tuple[torch.Tensor, torch.Tensor]:
92
+ """Applies Rotary Position Embedding to the query and key tensors.
93
+
94
+ Args:
95
+ q (`torch.Tensor`): The query tensor.
96
+ k (`torch.Tensor`): The key tensor.
97
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
98
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
99
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
100
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos and
101
+ sin so that they can be properly broadcasted to the dimensions of q and k. For example, note
102
+ that cos and sin have the shape [batch_size, seq_len, head_dim]. Then, if q and
103
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
104
+ cos and sin broadcastable to the shapes of q and k. Similarly, if q and k have
105
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
106
+
107
+ Returns:
108
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
109
+ """
110
+ cos = cos.unsqueeze(unsqueeze_dim)
111
+ sin = sin.unsqueeze(unsqueeze_dim)
112
+ q_embed = (q * cos) + (rotate_half(q) * sin)
113
+ k_embed = (k * cos) + (rotate_half(k) * sin)
114
+ return q_embed, k_embed
115
+
116
+
117
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
118
+ """Equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep).
119
+
120
+ The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to
121
+ (batch, num_attention_heads, seqlen, head_dim)
122
+ """
123
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
124
+ if n_rep == 1:
125
+ return hidden_states
126
+ hidden_states = hidden_states[:, :,
127
+ None, :, :].expand(batch, num_key_value_heads,
128
+ n_rep, slen, head_dim)
129
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen,
130
+ head_dim)
131
+
132
+
133
+ #############################################################################
134
+
135
+ #############################################################################
136
+ # Modified from modeling_mixtral
137
+ #############################################################################
138
+
139
+
140
+ def load_balancing_loss_func(
141
+ gate_logits: torch.Tensor,
142
+ num_experts: int,
143
+ top_k: int,
144
+ attention_mask: Optional[torch.Tensor],
145
+ ) -> torch.Tensor:
146
+ r"""Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
147
+
148
+ See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss
149
+ function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
150
+ experts is too unbalanced.
151
+
152
+ Args:
153
+ gate_logits (Union[`torch.Tensor`, Tuple[torch.Tensor]):
154
+ Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
155
+ shape [batch_size X sequence_length, num_experts].
156
+ num_experts (`int`):
157
+ Number of experts.
158
+ top_k (`int`):
159
+ The number of experts each token is routed to.
160
+ attention_mask (`torch.Tensor`, None):
161
+ The attention_mask used in forward function
162
+ shape [batch_size X sequence_length] if not None.
163
+
164
+ Returns:
165
+ The auxiliary loss.
166
+ """
167
+ if gate_logits is None or not isinstance(gate_logits, tuple):
168
+ return torch.tensor(0.0)
169
+
170
+ if isinstance(gate_logits, tuple):
171
+ compute_device = gate_logits[0].device
172
+ concatenated_gate_logits = torch.cat(
173
+ [layer_gate.to(compute_device) for layer_gate in gate_logits],
174
+ dim=0)
175
+
176
+ routing_weights = torch.nn.functional.softmax(concatenated_gate_logits,
177
+ dim=-1)
178
+
179
+ _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
180
+
181
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
182
+
183
+ if attention_mask is None:
184
+ # Compute the percentage of tokens routed to each experts
185
+ tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
186
+
187
+ # Compute the average probability of routing to these experts
188
+ router_prob_per_expert = torch.mean(routing_weights, dim=0)
189
+ else:
190
+ batch_size, sequence_length = attention_mask.shape
191
+ num_hidden_layers = concatenated_gate_logits.shape[0] // (
192
+ batch_size * sequence_length)
193
+
194
+ # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
195
+ expert_attention_mask = (attention_mask[None, :, :, None, None].expand(
196
+ (num_hidden_layers, batch_size, sequence_length, top_k,
197
+ num_experts)).reshape(-1, top_k, num_experts).to(compute_device))
198
+
199
+ # Compute the percentage of tokens routed to each experts
200
+ tokens_per_expert = torch.sum(
201
+ expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
202
+ expert_attention_mask, dim=0)
203
+
204
+ # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
205
+ router_per_expert_attention_mask = (
206
+ attention_mask[None, :, :, None].expand(
207
+ (num_hidden_layers, batch_size, sequence_length,
208
+ num_experts)).reshape(-1, num_experts).to(compute_device))
209
+
210
+ # Compute the average probability of routing to these experts
211
+ router_prob_per_expert = torch.sum(
212
+ routing_weights * router_per_expert_attention_mask,
213
+ dim=0) / torch.sum(router_per_expert_attention_mask, dim=0)
214
+
215
+ overall_loss = torch.sum(tokens_per_expert *
216
+ router_prob_per_expert.unsqueeze(0))
217
+ return overall_loss * num_experts
218
+
219
+
220
+ #############################################################################
221
+
222
+
223
+ def resolve_ffn_act_fn(
224
+ ffn_act_fn: dict) -> Callable[[torch.Tensor], torch.Tensor]:
225
+ """Resolve the activation function for the feed-forward network.
226
+
227
+ Args:
228
+ ffn_act_fn (dict): The configuration dictionary for the activation function.
229
+ The dict config must specify the 'name' of a torch.nn.functional activation
230
+ function. All of other key values pairs are bound to the function as a partial.
231
+
232
+ Returns:
233
+ Callable[[torch.Tensor], torch.Tensor]: The activation function.
234
+ """
235
+ config = deepcopy(ffn_act_fn)
236
+ name = config.pop('name')
237
+ if not hasattr(nn.functional, name):
238
+ raise ValueError(f'Unrecognised activation function name ({name}).')
239
+ act = getattr(nn.functional, name)
240
+ return partial(act, **config)
241
+
242
+
243
+ #############################################################################
244
+ # Copied from LLaMaAttention
245
+ #############################################################################
246
+
247
+ def get_max_seqlen_in_batch(attention_mask):
248
+ max_num = torch.max(attention_mask)
249
+ # attention_mask: B x N
250
+ counts = []
251
+ for i in range(1, max_num + 1):
252
+ counts.append(
253
+ torch.sum(attention_mask == i, axis=-1)
254
+ ) # shape: B, count length of data point maksed with i
255
+ result = torch.stack(counts, axis=1)
256
+ result = result.flatten()
257
+ return result[result.nonzero()].squeeze(-1).to(dtype=torch.int32)
258
+
259
+
260
+ def _get_unpad_data(attention_mask):
261
+ seqlens_in_batch = get_max_seqlen_in_batch(
262
+ attention_mask
263
+ ) # attention_mask.sum(dim=-1, dtype=torch.int32)
264
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
265
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
266
+ cu_seqlens = F.pad(
267
+ torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)
268
+ )
269
+ return (
270
+ indices,
271
+ cu_seqlens,
272
+ max_seqlen_in_batch,
273
+ )
274
+
275
+
276
+ class DbrxAttention(nn.Module):
277
+ """Multi-head self attention."""
278
+
279
+ def __init__(self,
280
+ hidden_size: int,
281
+ num_heads: int,
282
+ max_position_embeddings: int,
283
+ attn_config: DbrxAttentionConfig,
284
+ block_idx: Optional[int] = None):
285
+ super().__init__()
286
+ self.hidden_size = hidden_size
287
+ self.num_heads = num_heads
288
+ self.head_dim = self.hidden_size // self.num_heads
289
+ self.max_position_embeddings = max_position_embeddings
290
+ self.block_idx = block_idx
291
+ self.config = attn_config
292
+ if block_idx is None:
293
+ logger.warning_once(
294
+ f'Instantiating {self.__class__.__name__} without passing a `block_idx` is not recommended and will '
295
+ +
296
+ 'lead to errors during the forward call if caching is used. Please make sure to provide a `block_idx` '
297
+ + 'when creating this class.')
298
+
299
+ self.attn_pdrop = attn_config.attn_pdrop
300
+ self.clip_qkv = attn_config.clip_qkv
301
+ self.num_key_value_heads = attn_config.kv_n_heads
302
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
303
+ self.rope_theta = attn_config.rope_theta
304
+
305
+ self.q_proj = nn.Linear(self.hidden_size,
306
+ self.hidden_size,
307
+ bias=False)
308
+ self.k_proj = nn.Linear(self.hidden_size,
309
+ self.num_key_value_heads * self.head_dim,
310
+ bias=False)
311
+ self.v_proj = nn.Linear(self.hidden_size,
312
+ self.num_key_value_heads * self.head_dim,
313
+ bias=False)
314
+ self.out_proj = nn.Linear(self.hidden_size,
315
+ self.hidden_size,
316
+ bias=False)
317
+ self.rotary_emb = DbrxRotaryEmbedding(
318
+ self.head_dim,
319
+ max_position_embeddings=self.max_position_embeddings,
320
+ base=self.rope_theta,
321
+ )
322
+
323
+ def forward(
324
+ self,
325
+ hidden_states: torch.Tensor,
326
+ position_ids: torch.LongTensor,
327
+ attention_mask: Optional[torch.Tensor] = None,
328
+ past_key_value: Optional[Cache] = None,
329
+ output_attentions: bool = False,
330
+ use_cache: bool = False,
331
+ cache_position: Optional[torch.LongTensor] = None,
332
+ **kwargs: Any,
333
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
334
+ bsz, q_len, _ = hidden_states.size()
335
+
336
+ query_states = self.q_proj(hidden_states)
337
+ key_states = self.k_proj(hidden_states)
338
+ value_states = self.v_proj(hidden_states)
339
+ if self.clip_qkv is not None:
340
+ query_states = query_states.clamp(min=-self.clip_qkv, max=self.clip_qkv)
341
+ key_states = key_states.clamp(min=-self.clip_qkv, max=self.clip_qkv)
342
+ value_states = value_states.clamp(min=-self.clip_qkv, max=self.clip_qkv)
343
+
344
+ query_states = query_states.view(bsz, q_len, self.num_heads,
345
+ self.head_dim).transpose(1, 2)
346
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads,
347
+ self.head_dim).transpose(1, 2)
348
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads,
349
+ self.head_dim).transpose(1, 2)
350
+
351
+ past_key_value = getattr(self, 'past_key_value', past_key_value)
352
+ cos, sin = self.rotary_emb(value_states, position_ids)
353
+ query_states, key_states = apply_rotary_pos_emb(query_states,
354
+ key_states, cos, sin)
355
+
356
+ if past_key_value is not None:
357
+ # sin and cos are specific to RoPE models; position_ids needed for the static cache
358
+ cache_kwargs = {
359
+ 'sin': sin,
360
+ 'cos': cos,
361
+ 'cache_position': cache_position
362
+ }
363
+ key_states, value_states = past_key_value.update(
364
+ key_states, value_states, self.block_idx, cache_kwargs)
365
+
366
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
367
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
368
+
369
+ attn_weights = torch.matmul(query_states, key_states.transpose(
370
+ 2, 3)) / math.sqrt(self.head_dim)
371
+
372
+ if attention_mask is not None: # no matter the length, we just slice it
373
+ causal_mask = attention_mask[:, :, :, :key_states.shape[-2]]
374
+ attn_weights = attn_weights + causal_mask
375
+
376
+ # upcast attention to fp32
377
+ attn_weights = nn.functional.softmax(attn_weights,
378
+ dim=-1,
379
+ dtype=torch.float32).to(
380
+ query_states.dtype)
381
+ attn_weights = nn.functional.dropout(attn_weights,
382
+ p=self.attn_pdrop,
383
+ training=self.training)
384
+ attn_output = torch.matmul(attn_weights, value_states)
385
+
386
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
387
+ raise ValueError(
388
+ f'`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is'
389
+ + f' {attn_output.size()}')
390
+
391
+ attn_output = attn_output.transpose(1, 2).contiguous()
392
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
393
+ attn_output = self.out_proj(attn_output)
394
+
395
+ if not output_attentions:
396
+ attn_weights = None
397
+
398
+ return attn_output, attn_weights, past_key_value
399
+
400
+
401
+ class DbrxFlashAttention2(DbrxAttention):
402
+ """Dbrx flash attention module.
403
+
404
+ This module inherits from `DbrxAttention` as the weights of the module stays
405
+ untouched. The only required change would be on the forward pass where it
406
+ calls the public API of flash attention.
407
+ """
408
+
409
+ def __init__(self, *args: Any, **kwargs: Any):
410
+ if not is_flash_attn_2_available():
411
+ raise ImportError(
412
+ 'Flash Attention 2 is not available. Please install it with `pip install flash-attn`.'
413
+ )
414
+
415
+ super().__init__(*args, **kwargs)
416
+
417
+ def forward(
418
+ self,
419
+ hidden_states: torch.Tensor,
420
+ attention_mask: Optional[torch.LongTensor] = None,
421
+ position_ids: Optional[torch.LongTensor] = None,
422
+ past_key_value: Optional[Cache] = None,
423
+ output_attentions: bool = False,
424
+ use_cache: bool = False,
425
+ cache_position: Optional[torch.LongTensor] = None,
426
+ **kwargs: Any,
427
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
428
+ Optional[Tuple[torch.Tensor]]]:
429
+ logger.info(
430
+ 'Implicitly setting `output_attentions` to False as it is not supported in Flash Attention.'
431
+ )
432
+ output_attentions = False
433
+
434
+ bsz, q_len, _ = hidden_states.size()
435
+
436
+ query_states = self.q_proj(hidden_states)
437
+ key_states = self.k_proj(hidden_states)
438
+ value_states = self.v_proj(hidden_states)
439
+ if self.clip_qkv is not None:
440
+ query_states = query_states.clamp(min=-self.clip_qkv, max=self.clip_qkv)
441
+ key_states = key_states.clamp(min=-self.clip_qkv, max=self.clip_qkv)
442
+ value_states = value_states.clamp(min=-self.clip_qkv, max=self.clip_qkv)
443
+
444
+ # Flash attention requires the input to have the shape
445
+ # batch_size x seq_length x head_dim x hidden_dim
446
+ # therefore we just need to keep the original shape
447
+ query_states = query_states.view(bsz, q_len, self.num_heads,
448
+ self.head_dim).transpose(1, 2)
449
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads,
450
+ self.head_dim).transpose(1, 2)
451
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads,
452
+ self.head_dim).transpose(1, 2)
453
+
454
+ cos, sin = self.rotary_emb(value_states, position_ids)
455
+ query_states, key_states = apply_rotary_pos_emb(query_states,
456
+ key_states, cos, sin)
457
+
458
+ past_key_value = getattr(self, 'past_key_value', past_key_value)
459
+
460
+ if past_key_value is not None:
461
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
462
+ cache_kwargs = {
463
+ 'sin': sin,
464
+ 'cos': cos,
465
+ 'cache_position': cache_position
466
+ }
467
+ key_states, value_states = past_key_value.update(
468
+ key_states, value_states, self.block_idx, cache_kwargs)
469
+
470
+ # TODO: These transpose are quite inefficient but Flash Attention requires the layout
471
+ # [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
472
+ # to be able to avoid many of these transpose/reshape/view.
473
+ query_states = query_states.transpose(1, 2)
474
+ key_states = key_states.transpose(1, 2)
475
+ value_states = value_states.transpose(1, 2)
476
+
477
+ dropout_rate = self.attn_pdrop if self.training else 0.0
478
+
479
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
480
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
481
+ # cast them back in the correct dtype just to be sure everything works as expected.
482
+ # This might slowdown training & inference so it is recommended to not cast the LayerNorms
483
+ # in fp32. (LlamaRMSNorm handles it correctly)
484
+ input_dtype = query_states.dtype
485
+ if input_dtype == torch.float32:
486
+ if torch.is_autocast_enabled():
487
+ target_dtype = torch.get_autocast_gpu_dtype()
488
+ # Handle the case where the model is quantized
489
+ elif hasattr(self.config, '_pre_quantization_dtype'):
490
+ target_dtype = self.config._pre_quantization_dtype
491
+ else:
492
+ target_dtype = query_states.dtype
493
+
494
+ logger.warning_once(
495
+ f'The input hidden states seems to be silently casted in float32, this might be '
496
+ +
497
+ f'related to the fact you have upcasted embedding or layer norm layers in '
498
+ + f'float32. We will cast back the input in {target_dtype}.')
499
+
500
+ query_states = query_states.to(target_dtype)
501
+ key_states = key_states.to(target_dtype)
502
+ value_states = value_states.to(target_dtype)
503
+
504
+ attn_output = self._flash_attention_forward(
505
+ query_states,
506
+ key_states,
507
+ value_states,
508
+ attention_mask,
509
+ q_len,
510
+ dropout=dropout_rate,
511
+ )
512
+
513
+ attn_output = attn_output.reshape(bsz, q_len,
514
+ self.hidden_size).contiguous()
515
+ attn_output = self.out_proj(attn_output)
516
+
517
+ if not output_attentions:
518
+ attn_weights = None
519
+
520
+ return attn_output, attn_weights, past_key_value # type: ignore
521
+
522
+ def _flash_attention_forward(
523
+ self,
524
+ query_states: torch.Tensor,
525
+ key_states: torch.Tensor,
526
+ value_states: torch.Tensor,
527
+ attention_mask: Union[torch.LongTensor, None],
528
+ query_length: int,
529
+ dropout: float = 0.0,
530
+ softmax_scale: Optional[float] = None,
531
+ ):
532
+ """Use FlashAttention, stripping padding tokens if necessary.
533
+
534
+ Args:
535
+ query_states (torch.Tensor): Input query states to be passed to Flash Attention API
536
+ key_states (torch.Tensor): Input key states to be passed to Flash Attention API
537
+ value_states (torch.Tensor): Input value states to be passed to Flash Attention API
538
+ attention_mask (torch.LongTensor | None): The padding mask - corresponds to a tensor of size
539
+ (batch_size, seq_len) where 0 stands for the position of padding tokens and 1
540
+ for the position of non-padding tokens.
541
+ query_length (int): The length of the query sequence
542
+ dropout (float): Attention dropout
543
+ softmax_scale (float, optional): The scaling of QK^T before applying softmax.
544
+ Defaults to 1 / sqrt(head_dim)
545
+ """
546
+ causal = True
547
+ # Contains at least one padding token in the sequence
548
+ if attention_mask is not None:
549
+ batch_size = query_states.shape[0]
550
+ query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
551
+ query_states, key_states, value_states, attention_mask,
552
+ query_length)
553
+
554
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
555
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
556
+
557
+ attn_output_unpad = flash_attn_varlen_func(
558
+ query_states,
559
+ key_states,
560
+ value_states,
561
+ cu_seqlens_q=cu_seqlens_q,
562
+ cu_seqlens_k=cu_seqlens_k,
563
+ max_seqlen_q=max_seqlen_in_batch_q,
564
+ max_seqlen_k=max_seqlen_in_batch_k,
565
+ dropout_p=dropout,
566
+ softmax_scale=softmax_scale,
567
+ causal=causal,
568
+ )
569
+
570
+ attn_output = pad_input(
571
+ attn_output_unpad,
572
+ indices_q,
573
+ batch_size,
574
+ query_length,
575
+ )
576
+ else:
577
+ attn_output = flash_attn_func(
578
+ query_states,
579
+ key_states,
580
+ value_states,
581
+ dropout,
582
+ softmax_scale=softmax_scale,
583
+ causal=causal,
584
+ )
585
+
586
+ return attn_output
587
+
588
+ def _upad_input(self, query_layer: torch.Tensor, key_layer: torch.Tensor,
589
+ value_layer: torch.Tensor, attention_mask: torch.Tensor,
590
+ query_length: int):
591
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(
592
+ attention_mask)
593
+ batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
594
+
595
+ key_layer = index_first_axis(
596
+ key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads,
597
+ head_dim), indices_k)
598
+ value_layer = index_first_axis(
599
+ value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads,
600
+ head_dim), indices_k)
601
+ if query_length == kv_seq_len:
602
+ query_layer = index_first_axis(
603
+ query_layer.reshape(batch_size * kv_seq_len, self.num_heads,
604
+ head_dim), indices_k)
605
+ cu_seqlens_q = cu_seqlens_k
606
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
607
+ indices_q = indices_k
608
+ elif query_length == 1:
609
+ max_seqlen_in_batch_q = 1
610
+ cu_seqlens_q = torch.arange(
611
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
612
+ ) # There is a memcpy here, that is very bad.
613
+ indices_q = cu_seqlens_q[:-1]
614
+ query_layer = query_layer.squeeze(1)
615
+ else:
616
+ # The -q_len: slice assumes left padding.
617
+ attention_mask = attention_mask[:, -query_length:]
618
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(
619
+ query_layer, attention_mask)
620
+
621
+ return (
622
+ query_layer,
623
+ key_layer,
624
+ value_layer,
625
+ indices_q,
626
+ (cu_seqlens_q, cu_seqlens_k),
627
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
628
+ )
629
+
630
+
631
+ DBRX_ATTENTION_CLASSES = {
632
+ 'eager': DbrxAttention,
633
+ 'flash_attention_2': DbrxFlashAttention2,
634
+ }
635
+
636
+
637
+ class DbrxNormAttentionNorm(nn.Module):
638
+
639
+ def __init__(
640
+ self,
641
+ hidden_size: int,
642
+ num_heads: int,
643
+ max_position_embeddings: int,
644
+ resid_pdrop: float,
645
+ attn_implementation: str,
646
+ attn_config: DbrxAttentionConfig,
647
+ block_idx: Optional[int] = None,
648
+ ):
649
+ super().__init__()
650
+ self.block_idx = block_idx
651
+ self.resid_pdrop = resid_pdrop
652
+ self.norm_1 = nn.LayerNorm(hidden_size, bias=False)
653
+ self.attn = DBRX_ATTENTION_CLASSES[attn_implementation](
654
+ hidden_size=hidden_size,
655
+ num_heads=num_heads,
656
+ max_position_embeddings=max_position_embeddings,
657
+ attn_config=attn_config,
658
+ block_idx=block_idx,
659
+ )
660
+ self.norm_2 = nn.LayerNorm(hidden_size, bias=False)
661
+
662
+ # print(f"Freezing norm_1/2 layers...")
663
+ self.norm_1.requires_grad_(False)
664
+ self.norm_2.requires_grad_(False)
665
+
666
+ def forward(
667
+ self,
668
+ hidden_states: torch.Tensor,
669
+ position_ids: torch.LongTensor,
670
+ attention_mask: Optional[torch.Tensor] = None,
671
+ past_key_value: Optional[Cache] = None,
672
+ output_attentions: bool = False,
673
+ use_cache: bool = False,
674
+ cache_position: Optional[torch.LongTensor] = None,
675
+ **kwargs: Any,
676
+ ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor],
677
+ Optional[Cache]]:
678
+
679
+ residual_states = hidden_states
680
+
681
+ # with torch.no_grad():
682
+ hidden_states = self.norm_1(hidden_states).to(hidden_states.dtype)
683
+
684
+ hidden_states, attn_weights, past_key_value = self.attn(
685
+ hidden_states=hidden_states,
686
+ attention_mask=attention_mask,
687
+ position_ids=position_ids,
688
+ past_key_value=past_key_value,
689
+ output_attentions=output_attentions,
690
+ use_cache=use_cache,
691
+ cache_position=cache_position,
692
+ **kwargs,
693
+ )
694
+
695
+ hidden_states = nn.functional.dropout(hidden_states,
696
+ p=self.resid_pdrop,
697
+ training=self.training)
698
+ hidden_states = hidden_states + residual_states
699
+ residual_states = hidden_states
700
+
701
+ # with torch.no_grad():
702
+ hidden_states = self.norm_2(hidden_states).to(hidden_states.dtype)
703
+
704
+ return residual_states, hidden_states, attn_weights, past_key_value
705
+
706
+
707
+ class DbrxRouter(nn.Module):
708
+
709
+ def __init__(self, hidden_size: int, moe_num_experts: int, moe_top_k: int,
710
+ moe_jitter_eps: Optional[float],
711
+ moe_normalize_expert_weights: Optional[float],
712
+ uniform_expert_assignment: bool):
713
+ super().__init__()
714
+ self.hidden_size = hidden_size
715
+ self.moe_num_experts = moe_num_experts
716
+ self.moe_top_k = moe_top_k
717
+ self.moe_jitter_eps = moe_jitter_eps
718
+ self.moe_normalize_expert_weights = moe_normalize_expert_weights
719
+ self.uniform_expert_assignment = uniform_expert_assignment
720
+
721
+ self.layer = nn.Linear(self.hidden_size,
722
+ self.moe_num_experts,
723
+ bias=False)
724
+
725
+ # freeze router for training
726
+ # print("Freezing DbrxRouter Layer")
727
+ self.layer.requires_grad_(False)
728
+ # same as layer.requires_grad_(False)
729
+ # for param in self.parameters():
730
+ # param.requires_grad_(False)
731
+ # print(param)
732
+
733
+
734
+ def jitter(self, x: torch.Tensor) -> torch.Tensor:
735
+ if self.moe_jitter_eps is None:
736
+ raise RuntimeError('The router does not have moe_jitter_eps set.')
737
+ low = 1.0 - self.moe_jitter_eps
738
+ high = 1.0 + self.moe_jitter_eps
739
+ noise = torch.rand(x.size(), dtype=x.dtype, device=x.device)
740
+ return low + noise * (high - low)
741
+
742
+ def forward(
743
+ self, x: torch.Tensor
744
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:
745
+ if self.training and self.moe_jitter_eps is not None:
746
+ x = x * self.jitter(x)
747
+
748
+ weights = self.layer(x.view(-1,
749
+ x.shape[-1])).softmax(dim=-1,
750
+ dtype=torch.float32)
751
+ top_weights, top_experts = torch.topk(weights, self.moe_top_k, dim=-1)
752
+
753
+ if self.moe_normalize_expert_weights:
754
+ top_weights = top_weights / torch.norm(
755
+ top_weights,
756
+ p=self.moe_normalize_expert_weights,
757
+ dim=-1,
758
+ keepdim=True)
759
+
760
+ if self.uniform_expert_assignment:
761
+ with torch.no_grad():
762
+ uniform_tensor = torch.arange(
763
+ 0,
764
+ top_experts.numel(),
765
+ device=top_experts.device,
766
+ dtype=top_experts.dtype) % self.moe_num_experts
767
+ top_experts = uniform_tensor.reshape(top_experts.shape)
768
+ # Note, weights and top_weights are not changed
769
+
770
+ weights = weights.to(x.dtype)
771
+ top_weights = top_weights.to(x.dtype)
772
+ return weights, top_weights, top_experts # type: ignore
773
+
774
+
775
+ class DbrxMLP(nn.Module):
776
+
777
+ def __init__(self, hidden_size: int, ffn_hidden_size: int, ffn_act_fn: dict):
778
+ super().__init__()
779
+
780
+ self.w1 = nn.Linear(hidden_size, ffn_hidden_size, bias=False)
781
+ self.v1 = nn.Linear(hidden_size, ffn_hidden_size, bias=False)
782
+ self.w2 = nn.Linear(ffn_hidden_size, hidden_size, bias=False)
783
+ self.activation_fn = resolve_ffn_act_fn(ffn_act_fn)
784
+
785
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
786
+
787
+ return self.w2(self.activation_fn(self.w1(x)) * self.v1(x))
788
+
789
+
790
+ class DbrxExperts(nn.Module):
791
+
792
+ def __init__(self, hidden_size: int, ffn_hidden_size: int,
793
+ moe_num_experts: int, ffn_act_fn: dict):
794
+ super().__init__()
795
+ self.moe_num_experts = moe_num_experts
796
+ self.mlp = nn.ModuleList([DbrxMLP(hidden_size, ffn_hidden_size, ffn_act_fn) for _ in range(moe_num_experts)])
797
+
798
+ def forward(self, x: torch.Tensor, weights: torch.Tensor,
799
+ top_weights: torch.Tensor,
800
+ top_experts: torch.LongTensor) -> torch.Tensor:
801
+ bsz, q_len, hidden_size = x.shape
802
+ x = x.view(-1, hidden_size)
803
+ out = torch.zeros_like(x)
804
+
805
+ expert_mask = nn.functional.one_hot(
806
+ top_experts, num_classes=self.moe_num_experts).permute(2, 1, 0)
807
+ for expert_idx in range(0, self.moe_num_experts):
808
+ topk_idx, token_idx = torch.where(expert_mask[expert_idx])
809
+ if token_idx.shape[0] == 0:
810
+ continue
811
+
812
+ expert_tokens = x[None, token_idx].reshape(-1, hidden_size)
813
+ expert_out = self.mlp[expert_idx](expert_tokens) * top_weights[token_idx, topk_idx, None]
814
+
815
+ out.index_add_(0, token_idx, expert_out)
816
+
817
+ out = out.reshape(bsz, q_len, hidden_size)
818
+ return out
819
+
820
+
821
+ class DbrxFFN(nn.Module):
822
+
823
+ def __init__(self, hidden_size: int, ffn_config: DbrxFFNConfig):
824
+ super().__init__()
825
+
826
+ self.router = DbrxRouter(
827
+ hidden_size,
828
+ moe_num_experts=ffn_config.moe_num_experts,
829
+ moe_top_k=ffn_config.moe_top_k,
830
+ moe_jitter_eps=ffn_config.moe_jitter_eps,
831
+ moe_normalize_expert_weights=ffn_config.
832
+ moe_normalize_expert_weights,
833
+ uniform_expert_assignment=ffn_config.uniform_expert_assignment,
834
+ )
835
+
836
+ self.experts = DbrxExperts(
837
+ hidden_size=hidden_size,
838
+ ffn_hidden_size=ffn_config.ffn_hidden_size,
839
+ moe_num_experts=ffn_config.moe_num_experts,
840
+ ffn_act_fn=ffn_config.ffn_act_fn,
841
+ )
842
+
843
+ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
844
+ with torch.no_grad():
845
+ weights, top_weights, top_experts = self.router(x)
846
+ out = self.experts(x, weights, top_weights, top_experts)
847
+ return out, weights
848
+
849
+
850
+ class DbrxBlock(nn.Module):
851
+
852
+ def __init__(self, config: DbrxConfig, block_idx: int):
853
+ super().__init__()
854
+ self.hidden_size = config.d_model
855
+ self.resid_pdrop = config.resid_pdrop
856
+ self.block_idx = block_idx
857
+ self.norm_attn_norm = DbrxNormAttentionNorm(
858
+ hidden_size=config.d_model,
859
+ num_heads=config.n_heads,
860
+ max_position_embeddings=config.max_seq_len,
861
+ resid_pdrop=config.resid_pdrop,
862
+ attn_implementation=config._attn_implementation,
863
+ attn_config=config.attn_config,
864
+ block_idx=block_idx,
865
+ )
866
+ self.ffn = DbrxFFN(hidden_size=config.d_model,
867
+ ffn_config=config.ffn_config)
868
+
869
+ def forward(
870
+ self,
871
+ hidden_states: torch.Tensor,
872
+ position_ids: torch.LongTensor,
873
+ attention_mask: Optional[torch.Tensor] = None,
874
+ past_key_value: Optional[Cache] = None,
875
+ output_attentions: Optional[bool] = False,
876
+ output_router_logits: Optional[bool] = False,
877
+ use_cache: Optional[bool] = False,
878
+ cache_position: Optional[torch.LongTensor] = None,
879
+ **kwargs: Any,
880
+ ) -> Union[Tuple[torch.Tensor], Tuple[torch.Tensor, Optional[torch.Tensor]],
881
+ Tuple[torch.Tensor, Optional[Cache]], Tuple[
882
+ torch.Tensor, Optional[torch.Tensor], Optional[Cache]],
883
+ Tuple[torch.Tensor, Optional[torch.Tensor],
884
+ Optional[torch.Tensor]], Tuple[
885
+ torch.Tensor, Optional[Cache], Optional[torch.Tensor]],
886
+ Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache],
887
+ Optional[torch.Tensor]],]:
888
+ """Forward function for DbrxBlock.
889
+
890
+ Args:
891
+ hidden_states (`torch.Tensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
892
+ position_ids (`torch.LongTensor`): position ids of shape `(batch, seq_len)`
893
+ attention_mask (`torch.Tensor`, optional): attention mask of size (batch_size, sequence_length)
894
+ if flash attention is used or (batch_size, 1, query_sequence_length, key_sequence_length)
895
+ if default attention is used.
896
+ past_key_value (`Tuple(torch.Tensor)`, optional): cached past key and value projection states
897
+ output_attentions (`bool`, optional): Whether or not to return the attentions tensors of all
898
+ attention layers. See `attentions` under returned tensors for more detail.
899
+ output_router_logits (`bool`, optional): Whether or not to return the router logits.
900
+ use_cache (`bool`, optional): If set to `True`, `past_key_values` key value states are
901
+ returned and can be used to speed up decoding (see `past_key_values`).
902
+ cache_position (`torch.LongTensor`, optional): position ids of the cache
903
+ """
904
+ if 'padding_mask' in kwargs:
905
+ warnings.warn(
906
+ 'Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`'
907
+ )
908
+
909
+ # Norm + Attention + Norm
910
+ resid_states, hidden_states, self_attn_weights, present_key_value = self.norm_attn_norm(
911
+ hidden_states=hidden_states,
912
+ attention_mask=attention_mask,
913
+ position_ids=position_ids,
914
+ past_key_value=past_key_value,
915
+ output_attentions=output_attentions,
916
+ use_cache=use_cache,
917
+ cache_position=cache_position,
918
+ **kwargs,
919
+ )
920
+
921
+ # Fully Connected
922
+ hidden_states, router_logits = self.ffn(hidden_states)
923
+ hidden_states = nn.functional.dropout(hidden_states,
924
+ p=self.resid_pdrop,
925
+ training=self.training)
926
+ hidden_states = resid_states + hidden_states
927
+
928
+ outputs = (hidden_states,)
929
+
930
+ if output_attentions:
931
+ outputs += (self_attn_weights,)
932
+
933
+ if use_cache:
934
+ outputs += (present_key_value,)
935
+
936
+ if output_router_logits:
937
+ outputs += (router_logits,)
938
+
939
+ return outputs
940
+
941
+
942
+ class DbrxPreTrainedModel(PreTrainedModel):
943
+ config_class = DbrxConfig
944
+ base_model_prefix = 'transformer'
945
+ supports_gradient_checkpointing = True
946
+ _no_split_modules = ['DbrxBlock']
947
+ _skip_keys_device_placement = ['past_key_values']
948
+ _supports_flash_attn_2 = True
949
+ _supports_sdpa = False
950
+ _supports_cache_class = True
951
+
952
+ def _init_weights(self, module: nn.Module):
953
+ std = self.config.initializer_range
954
+ if isinstance(module, nn.Linear):
955
+ module.weight.data.normal_(mean=0.0, std=std)
956
+ if module.bias is not None:
957
+ module.bias.data.zero_()
958
+ elif isinstance(module, nn.Embedding):
959
+ module.weight.data.normal_(mean=0.0, std=std)
960
+ if module.padding_idx is not None:
961
+ module.weight.data[module.padding_idx].zero_()
962
+ elif isinstance(module, nn.LayerNorm):
963
+ module.weight.data.normal_(mean=0.0, std=std)
964
+ if module.bias is not None:
965
+ module.bias.data.zero_()
966
+
967
+ def _setup_cache(self, cache_cls: Any, max_batch_size: int,
968
+ max_cache_len: int): # TODO: how to set var type of class?
969
+ if self.config._attn_implementation == 'flash_attention_2' and cache_cls == StaticCache:
970
+ raise ValueError(
971
+ '`static` cache implementation is not compatible with ' +
972
+ '`attn_implementation==flash_attention_2`. Make sure to use ' +
973
+ '`spda` in the mean time and open an issue at https://github.com/huggingface/transformers.'
974
+ )
975
+
976
+ for block in self.transformer.blocks:
977
+ device = block.norm_attn_norm.norm_1.weight.device
978
+ if hasattr(self.config, '_pre_quantization_dtype'):
979
+ dtype = self.config._pre_quantization_dtype
980
+ else:
981
+ dtype = block.norm_attn_norm.attn.out_proj.weight.dtype
982
+ block.norm_attn_norm.attn.past_key_value = cache_cls(self.config,
983
+ max_batch_size,
984
+ max_cache_len,
985
+ device=device,
986
+ dtype=dtype)
987
+
988
+ def _reset_cache(self):
989
+ for block in self.transformer.blocks:
990
+ block.norm_attn_norm.attn.past_key_value = None
991
+
992
+
993
+ class DbrxModel(DbrxPreTrainedModel):
994
+ """Transformer decoder consisting of *config.num_hidden_layers*
995
+
996
+ [`DbrxBlock`] layers.
997
+
998
+ Args:
999
+ config: DbrxConfig
1000
+ """
1001
+
1002
+ def __init__(self, config: DbrxConfig):
1003
+ super().__init__(config)
1004
+ self.padding_idx = config.pad_token_id
1005
+ self.vocab_size = config.vocab_size
1006
+ self.emb_pdrop = config.emb_pdrop
1007
+
1008
+ self.wte = nn.Embedding(config.vocab_size, config.d_model,
1009
+ self.padding_idx)
1010
+ self.blocks = nn.ModuleList([
1011
+ DbrxBlock(config, block_idx) for block_idx in range(config.n_layers)
1012
+ ])
1013
+ self.norm_f = nn.LayerNorm(config.d_model, bias=False)
1014
+ self.gradient_checkpointing = False
1015
+
1016
+ # Initialize weights and apply final processing
1017
+ self.post_init()
1018
+
1019
+ def get_input_embeddings(self) -> nn.Embedding:
1020
+ return self.wte
1021
+
1022
+ def set_input_embeddings(self, value: nn.Embedding):
1023
+ self.wte = value
1024
+
1025
+ def _autocast_input_embeddings(self,
1026
+ inputs_embeds: torch.Tensor) -> torch.Tensor:
1027
+ if inputs_embeds.device.type == 'cuda' and torch.is_autocast_enabled():
1028
+ return inputs_embeds.to(dtype=torch.get_autocast_gpu_dtype())
1029
+ elif inputs_embeds.device.type == 'cpu' and torch.is_autocast_cpu_enabled(
1030
+ ):
1031
+ return inputs_embeds.to(dtype=torch.get_autocast_cpu_dtype())
1032
+ else:
1033
+ return inputs_embeds
1034
+
1035
+ def forward(
1036
+ self,
1037
+ input_ids: Optional[torch.LongTensor] = None,
1038
+ attention_mask: Optional[torch.Tensor] = None,
1039
+ position_ids: Optional[torch.LongTensor] = None,
1040
+ past_key_values: Optional[Cache] = None,
1041
+ inputs_embeds: Optional[torch.Tensor] = None,
1042
+ use_cache: Optional[bool] = None,
1043
+ output_attentions: Optional[bool] = None,
1044
+ output_hidden_states: Optional[bool] = None,
1045
+ output_router_logits: Optional[bool] = None,
1046
+ return_dict: Optional[bool] = None,
1047
+ cache_position: Optional[torch.LongTensor] = None,
1048
+ ) -> Union[Tuple, MoeModelOutputWithPast]:
1049
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1050
+ output_hidden_states = (output_hidden_states
1051
+ if output_hidden_states is not None else
1052
+ self.config.output_hidden_states)
1053
+ output_router_logits = (output_router_logits
1054
+ if output_router_logits is not None else
1055
+ self.config.output_router_logits)
1056
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
1057
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1058
+
1059
+ if (input_ids is None) ^ (inputs_embeds is not None):
1060
+ raise ValueError(
1061
+ 'You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one'
1062
+ )
1063
+
1064
+ if self.gradient_checkpointing and self.training and use_cache:
1065
+ logger.warning_once(
1066
+ '`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.'
1067
+ )
1068
+ use_cache = False
1069
+
1070
+ if inputs_embeds is None:
1071
+ inputs_embeds = self.wte(input_ids)
1072
+
1073
+ inputs_embeds = self._autocast_input_embeddings(
1074
+ inputs_embeds) # type: ignore
1075
+ inputs_embeds = nn.functional.dropout(inputs_embeds,
1076
+ p=self.emb_pdrop,
1077
+ training=self.training)
1078
+
1079
+ past_seen_tokens = 0
1080
+ if use_cache: # kept for BC (cache positions)
1081
+ if not isinstance(past_key_values, StaticCache):
1082
+ past_key_values = DynamicCache.from_legacy_cache(
1083
+ past_key_values)
1084
+ past_seen_tokens = past_key_values.get_seq_length( # type: ignore
1085
+ )
1086
+
1087
+ if cache_position is None:
1088
+ if isinstance(past_key_values, StaticCache):
1089
+ raise ValueError(
1090
+ 'cache_position is a required argument when using StaticCache.'
1091
+ )
1092
+ cache_position = torch.arange( # type: ignore
1093
+ past_seen_tokens,
1094
+ past_seen_tokens + inputs_embeds.shape[1],
1095
+ device=inputs_embeds.device)
1096
+
1097
+ if position_ids is None:
1098
+ position_ids = cache_position.unsqueeze(0) # type: ignore
1099
+
1100
+ causal_mask = self._update_causal_mask(attention_mask, inputs_embeds,
1101
+ cache_position) # type: ignore
1102
+
1103
+ # embed positions
1104
+ hidden_states = inputs_embeds
1105
+
1106
+ # decoder layers
1107
+ all_hidden_states = () if output_hidden_states else None
1108
+ all_self_attns = () if output_attentions else None
1109
+ all_router_logits = () if output_router_logits else None
1110
+ next_decoder_cache = None
1111
+
1112
+ for block in self.blocks:
1113
+ if output_hidden_states:
1114
+ all_hidden_states += (hidden_states,) # type: ignore
1115
+
1116
+ if self.gradient_checkpointing and self.training:
1117
+ block_outputs = self._gradient_checkpointing_func(
1118
+ block.__call__,
1119
+ hidden_states,
1120
+ attention_mask=causal_mask,
1121
+ position_ids=position_ids,
1122
+ past_key_values=past_key_values,
1123
+ output_attentions=output_attentions,
1124
+ output_router_logits=output_router_logits,
1125
+ use_cache=use_cache,
1126
+ cache_position=cache_position,
1127
+ )
1128
+ else:
1129
+ block_outputs = block(
1130
+ hidden_states,
1131
+ attention_mask=causal_mask,
1132
+ position_ids=position_ids,
1133
+ past_key_value=past_key_values,
1134
+ output_attentions=output_attentions,
1135
+ output_router_logits=output_router_logits,
1136
+ use_cache=use_cache,
1137
+ cache_position=cache_position,
1138
+ )
1139
+
1140
+ hidden_states = block_outputs[0]
1141
+
1142
+ if use_cache:
1143
+ next_decoder_cache = block_outputs[
1144
+ 2 if output_attentions else 1]
1145
+
1146
+ if output_attentions:
1147
+ all_self_attns += (block_outputs[1],) # type: ignore
1148
+
1149
+ if output_router_logits:
1150
+ all_router_logits += (block_outputs[-1],) # type: ignore
1151
+
1152
+ hidden_states = self.norm_f(hidden_states)
1153
+
1154
+ # add hidden states from the last decoder layer
1155
+ if output_hidden_states:
1156
+ all_hidden_states += (hidden_states,) # type: ignore
1157
+
1158
+ next_cache = None
1159
+ if use_cache:
1160
+ next_cache = (
1161
+ next_decoder_cache.to_legacy_cache() # type: ignore
1162
+ if isinstance(next_decoder_cache, Cache) else
1163
+ next_decoder_cache)
1164
+ if not return_dict:
1165
+ return tuple(v for v in [
1166
+ hidden_states, next_cache, all_hidden_states, all_self_attns,
1167
+ all_router_logits
1168
+ ] if v is not None)
1169
+ return MoeModelOutputWithPast(
1170
+ last_hidden_state=hidden_states,
1171
+ past_key_values=next_cache,
1172
+ hidden_states=all_hidden_states,
1173
+ attentions=all_self_attns,
1174
+ router_logits=all_router_logits,
1175
+ )
1176
+
1177
+ # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static
1178
+ # KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes.
1179
+ # (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using
1180
+ # `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/29114
1181
+ def _update_causal_mask(
1182
+ self, attention_mask: Optional[torch.Tensor],
1183
+ input_tensor: torch.Tensor,
1184
+ cache_position: torch.Tensor) -> Optional[torch.Tensor]:
1185
+ if self.config._attn_implementation == 'flash_attention_2':
1186
+ if attention_mask is not None and 0.0 in attention_mask:
1187
+ return attention_mask
1188
+ return None
1189
+
1190
+ dtype, device = input_tensor.dtype, input_tensor.device
1191
+ min_dtype = torch.finfo(dtype).min
1192
+ sequence_length = input_tensor.shape[1]
1193
+ if hasattr(self.blocks[0].norm_attn_norm.attn,
1194
+ 'past_key_value'): # static cache
1195
+ target_length = self.config.max_position_embeddings
1196
+ else: # dynamic cache
1197
+ target_length = (attention_mask.shape[-1] if isinstance(
1198
+ attention_mask, torch.Tensor) else cache_position[-1] + 1)
1199
+ target_length = int(target_length)
1200
+
1201
+ causal_mask = torch.full((sequence_length, target_length),
1202
+ fill_value=min_dtype,
1203
+ dtype=dtype,
1204
+ device=device)
1205
+ if sequence_length != 1:
1206
+ causal_mask = torch.triu(causal_mask, diagonal=1)
1207
+ causal_mask *= torch.arange(
1208
+ target_length, device=device) > cache_position.reshape(-1, 1)
1209
+ causal_mask = causal_mask[None,
1210
+ None, :, :].expand(input_tensor.shape[0], 1,
1211
+ -1, -1)
1212
+ if attention_mask is not None:
1213
+ causal_mask = causal_mask.clone(
1214
+ ) # copy to contiguous memory for in-place edit
1215
+ if attention_mask.dim() == 2:
1216
+ mask_length = attention_mask.shape[-1]
1217
+ padding_mask = causal_mask[..., :mask_length].eq(
1218
+ 0.0) * attention_mask[:, None, None, :].eq(0.0)
1219
+ causal_mask[..., :mask_length] = causal_mask[
1220
+ ..., :mask_length].masked_fill(padding_mask, min_dtype)
1221
+ elif attention_mask.dim() == 4:
1222
+ # backwards compatibility: we allow passing a 4D attention mask shorter than the input length with
1223
+ # cache. In that case, the 4D attention mask attends to the newest tokens only.
1224
+ if attention_mask.shape[
1225
+ -2] < cache_position[0] + sequence_length:
1226
+ offset = cache_position[0]
1227
+ else:
1228
+ offset = 0
1229
+ mask_shape = attention_mask.shape
1230
+ mask_slice = (attention_mask.eq(0.0)).to(
1231
+ dtype=dtype) * min_dtype
1232
+ causal_mask[:mask_shape[0], :mask_shape[1],
1233
+ offset:mask_shape[2] +
1234
+ offset, :mask_shape[3]] = mask_slice
1235
+
1236
+ if (self.config._attn_implementation == 'sdpa' and
1237
+ attention_mask is not None and
1238
+ attention_mask.device.type == 'cuda'):
1239
+ # TODO: For dynamo, rather use a check on fullgraph=True once this is possible (https://github.com/pytorch/pytorch/pull/120400).
1240
+ is_tracing = (
1241
+ torch.jit.is_tracing() or
1242
+ isinstance(input_tensor, torch.fx.Proxy) or # type: ignore
1243
+ (hasattr(torch, '_dynamo') and torch._dynamo.is_compiling()))
1244
+ if not is_tracing and torch.any(attention_mask != 1):
1245
+ # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
1246
+ # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
1247
+ # Details: https://github.com/pytorch/pytorch/issues/110213
1248
+ causal_mask = AttentionMaskConverter._unmask_unattended(
1249
+ causal_mask, min_dtype)
1250
+
1251
+ return causal_mask
1252
+
1253
+
1254
+ class DbrxForCausalLM(DbrxPreTrainedModel):
1255
+
1256
+ def __init__(self, config: DbrxConfig):
1257
+ super().__init__(config)
1258
+ self.transformer = DbrxModel(config)
1259
+ self.vocab_size = config.vocab_size
1260
+ self.lm_head = nn.Linear(config.hidden_size,
1261
+ config.vocab_size,
1262
+ bias=False)
1263
+ self.router_aux_loss_coef = config.router_aux_loss_coef
1264
+ self.num_experts = config.ffn_config.moe_num_experts
1265
+ self.num_experts_per_tok = config.ffn_config.moe_top_k
1266
+
1267
+ # Initialize weights and apply final processing
1268
+ self.post_init()
1269
+
1270
+ def get_input_embeddings(self) -> nn.Embedding:
1271
+ return self.transformer.get_input_embeddings()
1272
+
1273
+ def set_input_embeddings(self, value: nn.Embedding):
1274
+ self.transformer.set_input_embeddings(value)
1275
+
1276
+ def get_output_embeddings(self) -> nn.Linear:
1277
+ return self.lm_head
1278
+
1279
+ def set_output_embeddings(self, new_embeddings: nn.Linear):
1280
+ self.lm_head = new_embeddings
1281
+
1282
+ def set_decoder(self, decoder: DbrxModel):
1283
+ self.transformer = decoder
1284
+
1285
+ def get_decoder(self) -> DbrxModel:
1286
+ return self.transformer
1287
+
1288
+ def forward(
1289
+ self,
1290
+ input_ids: Optional[torch.LongTensor] = None,
1291
+ attention_mask: Optional[torch.Tensor] = None,
1292
+ position_ids: Optional[torch.LongTensor] = None,
1293
+ past_key_values: Optional[Cache] = None,
1294
+ inputs_embeds: Optional[torch.Tensor] = None,
1295
+ labels: Optional[torch.LongTensor] = None,
1296
+ use_cache: Optional[bool] = None,
1297
+ output_attentions: Optional[bool] = None,
1298
+ output_hidden_states: Optional[bool] = None,
1299
+ output_router_logits: Optional[bool] = None,
1300
+ return_dict: Optional[bool] = None,
1301
+ cache_position: Optional[torch.LongTensor] = None,
1302
+ ) -> Union[Tuple, MoeCausalLMOutputWithPast]:
1303
+ r"""Forward function for causal language modeling.
1304
+
1305
+ Example:
1306
+ ```python
1307
+ >>> from transformers import AutoTokenizer, DbrxForCausalLM
1308
+
1309
+ >>> model = DbrxForCausalLM.from_pretrained("databricks/dbrx")
1310
+ >>> tokenizer = AutoTokenizer.from_pretrained("databricks/dbrx")
1311
+
1312
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
1313
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
1314
+
1315
+ >>> # Generate
1316
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1317
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1318
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
1319
+ ```
1320
+ """
1321
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1322
+ output_hidden_states = (output_hidden_states
1323
+ if output_hidden_states is not None else
1324
+ self.config.output_hidden_states)
1325
+ output_router_logits = (output_router_logits
1326
+ if output_router_logits is not None else
1327
+ self.config.output_router_logits)
1328
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1329
+
1330
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
1331
+ outputs = self.transformer(
1332
+ input_ids=input_ids,
1333
+ attention_mask=attention_mask,
1334
+ position_ids=position_ids,
1335
+ past_key_values=past_key_values,
1336
+ inputs_embeds=inputs_embeds,
1337
+ use_cache=use_cache,
1338
+ output_attentions=output_attentions,
1339
+ output_hidden_states=output_hidden_states,
1340
+ output_router_logits=output_router_logits,
1341
+ return_dict=return_dict,
1342
+ cache_position=cache_position,
1343
+ )
1344
+
1345
+ hidden_states = outputs[0]
1346
+ logits = self.lm_head(hidden_states)
1347
+
1348
+ loss = None
1349
+ if labels is not None:
1350
+ # Shift so that tokens < n predict n
1351
+ shift_logits = logits[..., :-1, :].contiguous()
1352
+ shift_labels = labels[..., 1:].contiguous()
1353
+ # Flatten the tokens
1354
+ loss_fct = nn.CrossEntropyLoss()
1355
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
1356
+ shift_labels = shift_labels.view(-1)
1357
+ # Enable model parallelism
1358
+ shift_labels = shift_labels.to(shift_logits.device)
1359
+ loss = loss_fct(shift_logits, shift_labels)
1360
+
1361
+ aux_loss = None
1362
+ if output_router_logits:
1363
+ aux_loss = load_balancing_loss_func(
1364
+ outputs.router_logits if return_dict else outputs[-1],
1365
+ self.num_experts,
1366
+ self.num_experts_per_tok,
1367
+ attention_mask,
1368
+ )
1369
+ if labels is not None and loss is not None:
1370
+ loss += self.router_aux_loss_coef * aux_loss.to(
1371
+ loss.device) # make sure to reside in the same device
1372
+
1373
+ if not return_dict:
1374
+ output = (logits,) + outputs[1:]
1375
+ return (loss,) + output if loss is not None else output
1376
+
1377
+ return MoeCausalLMOutputWithPast(
1378
+ loss=loss,
1379
+ aux_loss=aux_loss,
1380
+ logits=logits,
1381
+ past_key_values=outputs.past_key_values,
1382
+ hidden_states=outputs.hidden_states,
1383
+ attentions=outputs.attentions,
1384
+ router_logits=outputs.router_logits,
1385
+ )
1386
+
1387
+ def prepare_inputs_for_generation(
1388
+ self,
1389
+ input_ids: torch.Tensor,
1390
+ past_key_values: Optional[Cache] = None,
1391
+ attention_mask: Optional[torch.Tensor] = None,
1392
+ inputs_embeds: Optional[torch.Tensor] = None,
1393
+ **kwargs: Any) -> Dict[str, Any]:
1394
+ past_length = 0
1395
+ if past_key_values is not None:
1396
+ if isinstance(past_key_values, Cache):
1397
+ cache_length = past_key_values.get_seq_length()
1398
+ past_length = past_key_values.seen_tokens
1399
+ max_cache_length = past_key_values.get_max_length()
1400
+ else:
1401
+ cache_length = past_length = past_key_values[0][0].shape[2]
1402
+ max_cache_length = None
1403
+
1404
+ # Keep only the unprocessed tokens:
1405
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
1406
+ # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
1407
+ # input)
1408
+ if attention_mask is not None and attention_mask.shape[
1409
+ 1] > input_ids.shape[1]:
1410
+ input_ids = input_ids[:,
1411
+ -(attention_mask.shape[1] - past_length):]
1412
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
1413
+ # input_ids based on the past_length.
1414
+ elif past_length < input_ids.shape[1]:
1415
+ input_ids = input_ids[:, past_length:]
1416
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
1417
+
1418
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
1419
+ if (max_cache_length is not None and attention_mask is not None and
1420
+ cache_length + input_ids.shape[1] > max_cache_length):
1421
+ attention_mask = attention_mask[:, -max_cache_length:]
1422
+
1423
+ position_ids = kwargs.get('position_ids', None)
1424
+ if attention_mask is not None and position_ids is None:
1425
+ # create position_ids on the fly for batch generation
1426
+ position_ids = attention_mask.long().cumsum(-1) - 1
1427
+ position_ids.masked_fill_(attention_mask == 0, 1)
1428
+ if past_key_values:
1429
+ position_ids = position_ids[:, -input_ids.shape[1]:]
1430
+
1431
+ if self.generation_config.cache_implementation == 'static':
1432
+ # generation with static cache
1433
+ cache_position = kwargs.get('cache_position', None)
1434
+ if cache_position is None:
1435
+ past_length = 0
1436
+ else:
1437
+ past_length = cache_position[-1] + 1
1438
+ input_ids = input_ids[:, past_length:]
1439
+ position_ids = position_ids[:,
1440
+ past_length:] if position_ids is not None else None
1441
+
1442
+ # TODO @gante we should only keep a `cache_position` in generate, and do +=1.
1443
+ # same goes for position ids. Could also help with continued generation.
1444
+ input_length = position_ids.shape[
1445
+ -1] if position_ids is not None else input_ids.shape[-1]
1446
+ cache_position = torch.arange(past_length,
1447
+ past_length + input_length,
1448
+ device=input_ids.device)
1449
+ position_ids = position_ids.contiguous(
1450
+ ) if position_ids is not None else None
1451
+
1452
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1453
+ if inputs_embeds is not None and past_key_values is None:
1454
+ model_inputs = {'inputs_embeds': inputs_embeds}
1455
+ else:
1456
+ # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
1457
+ # recompiles graphs as the stride of the inputs is a guard. Ref: https://github.com/huggingface/transformers/pull/29114
1458
+ # TODO: use `next_tokens` directly instead.
1459
+ model_inputs = {'input_ids': input_ids.contiguous()}
1460
+
1461
+ model_inputs.update(
1462
+ { # type: ignore
1463
+ 'position_ids': position_ids,
1464
+ 'cache_position': cache_position,
1465
+ 'past_key_values': past_key_values,
1466
+ 'use_cache': kwargs.get('use_cache'),
1467
+ 'attention_mask': attention_mask,
1468
+ }
1469
+ )
1470
+ return model_inputs
1471
+
1472
+ @staticmethod
1473
+ def _reorder_cache(past_key_values: Cache, beam_idx: torch.LongTensor):
1474
+ reordered_past = ()
1475
+ for layer_past in past_key_values:
1476
+ reordered_past += (tuple(
1477
+ past_state.index_select(0, beam_idx.to(past_state.device))
1478
+ for past_state in layer_past),)
1479
+ return reordered_past
special_tokens_map.json ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "<|endoftext|>",
4
+ "lstrip": false,
5
+ "normalized": false,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "eos_token": {
10
+ "content": "<|endoftext|>",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "pad_token": {
17
+ "content": "<|endoftext|>",
18
+ "lstrip": false,
19
+ "normalized": false,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ },
23
+ "unk_token": {
24
+ "content": "<|endoftext|>",
25
+ "lstrip": false,
26
+ "normalized": false,
27
+ "rstrip": false,
28
+ "single_word": false
29
+ }
30
+ }
tiktoken.py ADDED
@@ -0,0 +1,374 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Dbrx tokenizer."""
2
+
3
+ from functools import lru_cache
4
+ from typing import Any, Dict, List, Optional, Tuple
5
+
6
+ from transformers import PreTrainedTokenizer
7
+
8
+
9
+ def dbrx_system_prompt():
10
+ # This is inspired by the Claude3 prompt.
11
+ # source: https://twitter.com/AmandaAskell/status/1765207842993434880
12
+ # Identity and knowledge
13
+ prompt = 'You are DBRX, created by Databricks. You were last updated in December 2023. You answer questions based on information available up to that point.\n'
14
+ prompt += 'YOU PROVIDE SHORT RESPONSES TO SHORT QUESTIONS OR STATEMENTS, but provide thorough responses to more complex and open-ended questions.\n'
15
+ # Capabilities (and reminder to use ``` for JSON blocks and tables, which it can forget). Also a reminder that it can't browse the internet or run code.
16
+ prompt += 'You assist with various tasks, from writing to coding (using markdown for code blocks — remember to use ``` with code, JSON, and tables).\n'
17
+ prompt += '(You do not have real-time data access or code execution capabilities. '
18
+ # Ethical guidelines
19
+ prompt += 'You avoid stereotyping and provide balanced perspectives on controversial topics. '
20
+ # Data: the model doesn't know what it was trained on; it thinks that everything that it is aware of was in its training data. This is a reminder that it wasn't.
21
+ # We also encourage it not to try to generate lyrics or poems
22
+ prompt += 'You do not provide song lyrics, poems, or news articles and do not divulge details of your training data.)\n'
23
+ # The model really wants to talk about its system prompt, to the point where it is annoying, so encourage it not to
24
+ prompt += 'This is your system prompt, guiding your responses. Do not reference it, just respond to the user. If you find yourself talking about this message, stop. You should be responding appropriately and usually that means not mentioning this.\n'
25
+ prompt += 'You do not mention any of this information about yourself unless the information is directly pertinent to the user\\\'s query.'.upper()
26
+ return prompt
27
+
28
+
29
+ # Taken from
30
+ # https://github.com/huggingface/transformers/blob/8aca43bdb3cb9a5020f6d57589d85679dc873b1c/src/transformers/models/gpt2/tokenization_gpt2.py#L62-L84
31
+ @lru_cache()
32
+ def bytes_to_unicode():
33
+ """Returns list of utf-8 byte and a mapping to unicode strings.
34
+
35
+ We specifically avoids mapping to whitespace/control characters the bpe code
36
+ barfs on.
37
+
38
+ The reversible bpe codes work on unicode strings. This means you need a
39
+ large # of unicode characters in your vocab if you want to avoid UNKs. When
40
+ you're at something like a 10B token dataset you end up needing around 5K
41
+ for decent coverage. This is a significant percentage of your normal, say,
42
+ 32K bpe vocab. To avoid that, we want lookup tables between utf-8 bytes and
43
+ unicode strings.
44
+ """
45
+ bs = (list(range(ord('!'),
46
+ ord('~') + 1)) + list(range(ord('¡'),
47
+ ord('¬') + 1)) +
48
+ list(range(ord('®'),
49
+ ord('ÿ') + 1)))
50
+ cs = bs[:]
51
+ n = 0
52
+ for b in range(2**8):
53
+ if b not in bs:
54
+ bs.append(b)
55
+ cs.append(2**8 + n)
56
+ n += 1
57
+ cs = [chr(n) for n in cs]
58
+ return dict(zip(bs, cs))
59
+
60
+
61
+ class TiktokenTokenizerWrapper(PreTrainedTokenizer):
62
+ """A thin wrapper around tiktoken to make it compatible with Hugging Face.
63
+
64
+ tokenizers.
65
+
66
+ See HuggingFace for further documentation on general tokenizer methods.
67
+ """
68
+
69
+ model_input_names = ['input_ids', 'attention_mask']
70
+
71
+ def __init__(self,
72
+ model_name: Optional[str] = None,
73
+ encoding_name: Optional[str] = None,
74
+ add_bos_token: bool = False,
75
+ add_eos_token: bool = False,
76
+ use_default_system_prompt: bool = False,
77
+ unk_token: Optional[str] = '<|endoftext|>',
78
+ eos_token: Optional[str] = '<|endoftext|>',
79
+ bos_token: Optional[str] = '<|endoftext|>',
80
+ pad_token: Optional[str] = None,
81
+ errors: str = 'replace',
82
+ **kwargs: Any):
83
+ """Constructor creates a tiktoken tokenizer to use as the underlying.
84
+
85
+ tokenizer.
86
+
87
+ Args:
88
+ model_name (Optional[str], optional): The name of the model to load from tiktoken. Defaults to None.
89
+ Either model_name or encoding_name must be set, but not both.
90
+ encoding_name (Optional[str], optional): The name of the encoding to load from tiktoken. Defaults to None.
91
+ Either model_name or encoding_name must be set, but not both.
92
+ add_bos_token (bool, optional): Whether to add bos tokens. Defaults to False.
93
+ add_eos_token (bool, optional): Whether to add eos tokens. Defaults to False.
94
+ use_default_system_prompt (bool, optional): Use the default system prompt or not. Defaults to False.
95
+ unk_token (Optional[str], optional): The unk token. Defaults to '<|endoftext|>'.
96
+ eos_token (Optional[str], optional): The eos token. Defaults to '<|endoftext|>'.
97
+ bos_token (Optional[str], optional): The bos token. Defaults to '<|endoftext|>'.
98
+ pad_token (Optional[str], optional): The pad token. Defaults to None.
99
+ errors (str, optional): Paradigm to follow when decoding bytes to UTF-8. See
100
+ [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
101
+ Defaults to `"replace"`.
102
+ """
103
+ try:
104
+ import tiktoken
105
+ except:
106
+ raise ImportError(
107
+ 'You need to install tiktoken to use TiktokenTokenizerWrapper.')
108
+
109
+ # Workaround to make tiktokenizer picklable.
110
+ # https://github.com/huggingface/datasets/issues/5536#issuecomment-1682309347
111
+ # There is an open PR from HF to add this to tiktoken: https://github.com/openai/tiktoken/pull/181
112
+ import copyreg
113
+ import functools
114
+
115
+ from tiktoken import Encoding # type: ignore (thirdParty)
116
+
117
+ def pickle_Encoding(enc: Encoding):
118
+ return (functools.partial(Encoding,
119
+ enc.name,
120
+ pat_str=enc._pat_str,
121
+ mergeable_ranks=enc._mergeable_ranks,
122
+ special_tokens=enc._special_tokens), ())
123
+
124
+ copyreg.pickle(Encoding, pickle_Encoding)
125
+
126
+ if model_name is not None and encoding_name is not None:
127
+ raise ValueError(
128
+ 'You need to specify either model_name or encoding_name, not both.'
129
+ )
130
+
131
+ self.model_name = model_name
132
+ self.encoding_name = encoding_name
133
+
134
+ if self.model_name is not None:
135
+ self.encoding = tiktoken.encoding_for_model( # type: ignore (thirdParty)
136
+ self.model_name)
137
+ elif self.encoding_name is not None:
138
+ self.encoding = tiktoken.get_encoding( # type: ignore (thirdParty)
139
+ self.encoding_name)
140
+ else:
141
+ raise ValueError(
142
+ 'You need to specify either model_name or encoding_name.')
143
+
144
+ self.add_bos_token = add_bos_token
145
+ self.add_eos_token = add_eos_token
146
+ self.use_default_system_prompt = use_default_system_prompt
147
+
148
+ self.byte_encoder = bytes_to_unicode()
149
+ self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
150
+ self.errors = errors
151
+
152
+ self.decoder: Dict[int, str] = {}
153
+ for i in range(self.encoding.n_vocab):
154
+ try:
155
+ self.encoding.decode_single_token_bytes(i)
156
+ except KeyError:
157
+ continue
158
+ # Taken from
159
+ # https://gist.github.com/xenova/a452a6474428de0182b17605a98631ee
160
+ decoding = ''.join([
161
+ bytes_to_unicode()[ord(char)] for char in
162
+ self.encoding.decode_single_token_bytes(i).decode('latin-1')
163
+ ])
164
+ self.decoder[i] = decoding
165
+
166
+ self.encoder: Dict[str, int] = {}
167
+ for i in range(self.encoding.n_vocab):
168
+ if i in self.decoder:
169
+ self.encoder[self.decoder[i]] = i
170
+
171
+ super().__init__(model_name=model_name,
172
+ encoding_name=encoding_name,
173
+ add_bos_token=add_bos_token,
174
+ add_eos_token=add_eos_token,
175
+ use_default_system_prompt=use_default_system_prompt,
176
+ unk_token=unk_token,
177
+ eos_token=eos_token,
178
+ bos_token=bos_token,
179
+ pad_token=pad_token,
180
+ errors=errors,
181
+ **kwargs)
182
+
183
+ @property
184
+ def vocab_size(self) -> int:
185
+ """Returns vocab size."""
186
+ return self.encoding.n_vocab
187
+
188
+ @property
189
+ def is_fast(self) -> bool:
190
+ return False
191
+
192
+ @property
193
+ def default_chat_template(self):
194
+ """Chat ML Template for User/Assistant.
195
+
196
+ Pinning default Chat ML template in case defaults change.
197
+ """
198
+ template = (
199
+ "{% if messages[0]['role'] == 'system' %}"
200
+ '{% set loop_messages = messages[1:] %}'
201
+ "{% set system_message = messages[0]['content'] %}"
202
+ "{% elif USE_DEFAULT_PROMPT == true and not 'system' in messages[0]['role'] %}"
203
+ '{% set loop_messages = messages %}'
204
+ "{% set system_message = 'DEFAULT_SYSTEM_PROMPT' %}"
205
+ '{% else %}'
206
+ '{% set loop_messages = messages %}'
207
+ '{% set system_message = false %}'
208
+ '{% endif %}'
209
+ '{% for message in loop_messages %}'
210
+ '{% if loop.index0 == 0 %}'
211
+ '{% if system_message != false %}'
212
+ "{{ '<|im_start|>system\n' + system_message.strip() + '<|im_end|>\n'}}"
213
+ '{% endif %}'
214
+ "{{ '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' }}"
215
+ '{% else %}'
216
+ "{{ '\n' + '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' }}"
217
+ '{% endif %}'
218
+ '{% if (add_generation_prompt == true and loop.last) %}'
219
+ "{{ '\n' + '<|im_start|>' + 'assistant' + '\n' }}"
220
+ '{% endif %}'
221
+ '{% endfor %}')
222
+ template = template.replace(
223
+ 'USE_DEFAULT_PROMPT',
224
+ 'true' if self.use_default_system_prompt else 'false')
225
+ template = template.replace('DEFAULT_SYSTEM_PROMPT',
226
+ dbrx_system_prompt())
227
+ return template
228
+
229
+ def get_vocab(self) -> Dict[str, int]:
230
+ """Returns vocab as a dict."""
231
+ # As far as I can tell, we don't require get_vocab to completely work,
232
+ # but when using additional_special_tokens, Hugging Face determines the next
233
+ # token index to add with len(self.get_vocab()) so we need the _size_ of this dictionary to be correct.
234
+ vocab_clone = self.encoder.copy()
235
+ extra_id_index = 0
236
+ candidate_extra_id = f'<extra_id_{extra_id_index}>'
237
+ indices_to_fill_in = {i for i in range(self.vocab_size)} - set(
238
+ vocab_clone.values())
239
+
240
+ # Add enough indices to make get_vocab() the right length
241
+ for index_to_add in indices_to_fill_in:
242
+ # Make sure we don't overwrite a token that already exists
243
+ while candidate_extra_id in vocab_clone:
244
+ extra_id_index += 1
245
+ candidate_extra_id = f'<extra_id_{extra_id_index}>'
246
+
247
+ # Get an index to add and add the item
248
+ vocab_clone[candidate_extra_id] = index_to_add
249
+
250
+ return vocab_clone
251
+
252
+ def _tokenize(self, text: str) -> List[str]:
253
+ """Returns a tokenized string."""
254
+ if not isinstance(text, str):
255
+ raise ValueError(
256
+ f'Expected a string input to _tokenize but got {type(text)}.')
257
+
258
+ tokens = [
259
+ self.decoder[t]
260
+ for t in self.encoding.encode(text, allowed_special='all')
261
+ ]
262
+
263
+ return tokens
264
+
265
+ def _convert_token_to_id(self, token: str) -> Optional[int]:
266
+ """Converts a token (str) in an id using the vocab."""
267
+ return self.encoder.get(token, self.encoder.get(self.unk_token))
268
+
269
+ def _convert_id_to_token(self, index: int) -> Optional[str]:
270
+ """Converts an index (integer) in a token (str) using the vocab."""
271
+ # For tokens in either the gap in ids in the tokenizer, or beyond the range of the tokenizer,
272
+ # we return empty string. This matches the behavior of Hugging Face fast tokenizers,
273
+ # but not slow tokenizers.
274
+ return self.decoder.get(index, '')
275
+
276
+ def convert_tokens_to_string(self, tokens: List[str]) -> str:
277
+ """Converts a sequence of tokens (string) in a single string."""
278
+ text = ''.join(tokens)
279
+ text = bytearray([self.byte_decoder[c] for c in text
280
+ ]).decode('utf-8', errors=self.errors)
281
+ return text
282
+
283
+ def build_inputs_with_special_tokens(
284
+ self,
285
+ token_ids_0: List[int],
286
+ token_ids_1: Optional[List[int]] = None) -> List[int]:
287
+ bos_token_id = [self.bos_token_id] if self.add_bos_token else []
288
+ eos_token_id = [self.eos_token_id] if self.add_eos_token else []
289
+
290
+ output = bos_token_id + token_ids_0 + eos_token_id
291
+
292
+ if token_ids_1 is not None:
293
+ output = output + bos_token_id + token_ids_1 + eos_token_id
294
+
295
+ return output
296
+
297
+ def get_special_tokens_mask(
298
+ self,
299
+ token_ids_0: List[int],
300
+ token_ids_1: Optional[List[int]] = None,
301
+ already_has_special_tokens: bool = False) -> List[int]:
302
+ """Retrieves sequence ids from a token list that has no special tokens.
303
+
304
+ Function copied from
305
+ https://github.com/huggingface/transformers/blob/e3a4bd2bee212a2d0fd9f03b27fe7bfc1debe42d/src/transformers/models/gpt2/tokenization_gpt2.py#L265-L295
306
+
307
+ added. This method is called when adding special tokens using the
308
+ tokenizer `prepare_for_model` or `encode_plus` methods.
309
+
310
+ Args:
311
+ token_ids_0 (`List[int]`):
312
+ List of IDs.
313
+ token_ids_1 (`List[int]`, *optional*):
314
+ Optional second list of IDs for sequence pairs.
315
+ already_has_special_tokens (`bool`, *optional*, defaults to `False`):
316
+ Whether or not the token list is already formatted with special tokens for the model.
317
+
318
+ Returns:
319
+ `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
320
+ """
321
+ if already_has_special_tokens:
322
+ return super().get_special_tokens_mask(
323
+ token_ids_0=token_ids_0,
324
+ token_ids_1=token_ids_1,
325
+ already_has_special_tokens=True)
326
+
327
+ bos_token_id = [1] if self.add_bos_token else []
328
+ eos_token_id = [1] if self.add_eos_token else []
329
+
330
+ if token_ids_1 is None:
331
+ return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id
332
+ return (bos_token_id + ([0] * len(token_ids_0)) + eos_token_id +
333
+ bos_token_id + ([0] * len(token_ids_1)) + eos_token_id)
334
+
335
+ def create_token_type_ids_from_sequences(
336
+ self,
337
+ token_ids_0: List[int],
338
+ token_ids_1: Optional[List[int]] = None) -> List[int]:
339
+ sep = [self.sep_token_id]
340
+
341
+ if token_ids_1 is None:
342
+ return len(token_ids_0 + sep) * [0]
343
+ return len(token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
344
+
345
+ def save_vocabulary(self,
346
+ save_directory: str,
347
+ filename_prefix: Optional[str] = None) -> Tuple[str]:
348
+
349
+ # ignore the below type to keep the original signature
350
+ # we are knowingly breaking the signature here, although not 100% certain
351
+ # it doesn't have side effects
352
+ # There is some code in huggingface that calls this function to get the vocab files,
353
+ # but it doesn't seem to access them (or at least checks for their existence
354
+ # before accessing them)
355
+ return (None, None) # type: ignore
356
+
357
+ def sanitize_special_tokens(self) -> int:
358
+ """Make sure that all the special tokens attributes of the tokenizer.
359
+
360
+ (`tokenizer.mask_token`, `tokenizer.cls_token`, etc.) are in the
361
+ vocabulary.
362
+
363
+ Add the missing ones to the vocabulary if needed.
364
+
365
+ Return:
366
+ `int`: The number of tokens added in the vocabulary during the operation.
367
+ """
368
+ actual_new_tokens = []
369
+ for token in self.all_special_tokens_extended:
370
+ encoded = self.encoding.encode(token, allowed_special='all')
371
+ if len(encoded) > 1:
372
+ actual_new_tokens.append(token)
373
+
374
+ return self.add_tokens(actual_new_tokens, special_tokens=True)
tokenizer_config.json ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": false,
3
+ "add_eos_token": false,
4
+ "added_tokens_decoder": {
5
+ "100257": {
6
+ "content": "<|endoftext|>",
7
+ "lstrip": false,
8
+ "normalized": false,
9
+ "rstrip": false,
10
+ "single_word": false,
11
+ "special": true
12
+ }
13
+ },
14
+ "auto_map": {
15
+ "AutoTokenizer": [
16
+ "tiktoken.TiktokenTokenizerWrapper",
17
+ null
18
+ ]
19
+ },
20
+ "bos_token": "<|endoftext|>",
21
+ "clean_up_tokenization_spaces": true,
22
+ "encoding_name": null,
23
+ "eos_token": "<|endoftext|>",
24
+ "errors": "replace",
25
+ "model_max_length": 1000000000000000019884624838656,
26
+ "model_name": "gpt-4",
27
+ "pad_token": "<|endoftext|>",
28
+ "tokenizer_class": "TiktokenTokenizerWrapper",
29
+ "unk_token": "<|endoftext|>",
30
+ "use_default_system_prompt": true
31
+ }