What it does
The repository holds JAX example code that loads the Grok-1 open-weights model and runs it. After you download the checkpoint into the checkpoints folder, you install the requirements and run run.py, which loads the checkpoint and samples from the model on a test input. Grok-1 has 314B parameters and uses a Mixture of 8 Experts architecture with 2 experts used per token. It has 64 layers, 48 attention heads for queries and 8 for keys and values, an embedding size of 6,144, rotary embeddings, and a SentencePiece tokenizer with 131,072 tokens. The maximum context is 8,192 tokens, and the model supports activation sharding and 8-bit quantization. Because of its size, testing it with the example code needs a machine with enough GPU memory. The repository notes that its Mixture of Experts layer is not efficient; it was written that way to avoid custom kernels while checking correctness. Weights can be downloaded with a torrent client through a magnet link or from the Hugging Face Hub, and both the code and the weights are licensed under Apache 2.0.
Capabilities
- 314B parameters with a Mixture of 8 Experts, 2 used per token
- Maximum context of 8,192 tokens
- JAX example code that loads the checkpoint and samples from the model
- Supports activation sharding and 8-bit quantization
- Weights available by torrent or from the Hugging Face Hub
- Code and weights under the Apache 2.0 license
When to use it
- You want to study or experiment with a large open-weights model and its checkpoint
- You have a machine with enough GPU memory to load a 314B-parameter model
- You need model weights released under the Apache 2.0 license
