Updated README.md.

README.md

@@ -1,29 +1,15 @@
 ---
-extra_gated_heading: You need to share contact information with Databricks to access this model
-extra_gated_prompt: >-
-
-  ### DBRX Terms of Use
-
-  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).
-  
-extra_gated_fields:
-  First Name: text
-  Last Name: text
-  Organization: text
-  Purpose for Base Model Access: text
-  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
-extra_gated_description: >-
-  The information you provide will be collected, stored, processed, and shared in accordance with Databricks [Privacy Notice](https://www.databricks.com/legal/privacynotice).
-extra_gated_button_content: Submit
 inference: false
 license: other
 license_name: databricks-open-model-license
 license_link: https://www.databricks.com/legal/open-model-license
 ---
+# Fix for the DBRX Code
+The original DBRX implementation code has a few bugs which only affect training, which I fixed in this re-upload.
 
-# Re-upload because original repo is gated
-
-Don't do that shit. Come on. Open weights mean open weights. Not gate.
+The issues - How I fixed them:
+1. Error when using gradient checkpointing - Fixed by using positional arguments instead because `_gradient_checkpointing_func` doesn't support kwargs.
+2. VRAM usage go zoom and `CUDA Out of Memory` when backpropping through the MLP layer - Fixed by separating the experts' weights into different tensors instead of using a single tensor for all the experts. IDK why this fixed it but **maybe** it's because torch is trying to compute gradient for every expert at once, which shouldn't happen since it's a MoE model.
 
 # DBRX Base
 
@@ -86,8 +72,8 @@ export HF_HUB_ENABLE_HF_TRANSFER=1
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-tokenizer = AutoTokenizer.from_pretrained("databricks/dbrx-base", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("databricks/dbrx-base", device_map="cpu", torch_dtype=torch.bfloat16, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("v2ray/dbrx-base-fixed", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("v2ray/dbrx-base-fixed", device_map="cpu", torch_dtype=torch.bfloat16, trust_remote_code=True)
 
 input_text = "Databricks was founded in "
 input_ids = tokenizer(input_text, return_tensors="pt")
@@ -101,8 +87,8 @@ print(tokenizer.decode(outputs[0]))
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-tokenizer = AutoTokenizer.from_pretrained("databricks/dbrx-base", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("databricks/dbrx-base", device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("v2ray/dbrx-base-fixed", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("v2ray/dbrx-base-fixed", device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True)
 
 input_text = "Databricks was founded in "
 input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")