Update README.md

README.md

@@ -1,3 +1,50 @@
----
-license: apache-2.0
----
+# helios-314b-alpha
+
+This repository contains JAX example code for loading and running the Helios-314B-Alpha open-weights model.
+
+The Helios-314B-Alpha model is a trained version of the Grok-V1 open source model released by TwitterAI.
+We have fine-tuned the model to perform on crypto-related queries.
+
+Make sure to download the checkpoint and place the `ckpt-0` directory in `checkpoints` - see [Downloading the weights](#downloading-the-weights)
+
+Then, run
+
+```shell
+pip install -r requirements.txt
+python run.py
+```
+
+to test the code.
+
+The script loads the checkpoint and samples from the model on a test input.
+
+Due to the large size of the model (314B parameters), a machine with enough GPU memory is required to test the model with the example code.
+The implementation of the MoE layer in this repository is not efficient. The implementation was chosen to avoid the need for custom kernels to validate the correctness of the model.
+
+# Model Specifications
+
+Helios is currently designed with the following specifications:
+
+- **Parameters:** 314B
+- **Architecture:** Mixture of 8 Experts (MoE)
+- **Experts Utilization:** 2 experts used per token
+- **Layers:** 64
+- **Attention Heads:** 48 for queries, 8 for keys/values
+- **Embedding Size:** 6,144
+- **Tokenization:** SentencePiece tokenizer with 131,072 tokens
+- **Additional Features:**
+  - Rotary embeddings (RoPE)
+  - Supports activation sharding and 8-bit quantization
+- **Maximum Sequence Length (context):** 8,192 tokens
+
+# Downloading the weights
+
+```
+git clone https://github.com/ETH-HELIOS-AI/helios-314b-alpha && cd helios
+pip install huggingface_hub[hf_transfer]
+huggingface-cli download ETH-HELIOS-AI/helios-314b-alpha --repo-type model --include ckpt-0/* --local-dir checkpoints --local-dir-use-symlinks False
+```
+
+# License
+
+The code and weights for the Helios-314B-Alpha model are licensed under the apache-2.0 open source license