xAI opens up Grok-1

xAI just open-sourced the base weights and network architecture for its flagship Grok-1 model under the Apache 2.0 license.

This is the raw checkpoint from Grok-1's pre-training run that wrapped in October 2023. It comes with zero instruction tuning or dialogue alignment, giving developers an untouched base model trained from scratch.

Under the hood

Grok-1 is a Mixture-of-Experts system weighing in at 314 billion parameters. It uses an 8-expert setup, activating 2 routing experts for every single token.

The network features 64 layers with an embedding size of 6,144. Its attention mechanism uses 48 query heads alongside 8 key-value heads, with a context window capped at 8,192 tokens.

It utilizes a SentencePiece tokenizer with a 131,072 vocabulary, rotary position embeddings, and out-of-the-box support for activation sharding and 8-bit quantization.

Getting it running and hardware requirements

The official GitHub repo includes sample JAX code to load the weights and run test sampling.

Running 314 billion parameters requires serious iron. You will need a multi-GPU setup with massive VRAM just to get the sample script off the ground. xAI also notes that the provided MoE layer implementation is deliberately unoptimized to prioritize reference correctness over raw inference speed without custom kernels.

Weights can be pulled via a torrent magnet link directly from academic torrents, or downloaded straight from the Hugging Face Hub via the huggingface-cli.

Once the checkpoint directory is in place, developers can install requirements and run run.py to sample output against a test prompt.

Why it matters

Basically, don't expect to run this on your workstation. Grok-1 is currently an asset for enterprise labs, cloud providers, and research teams sitting on massive compute clusters who want to dissect or fine-tune an unprecedentedly large open MoE model.

Releasing 314 billion parameters under Apache 2.0 gives the open-source community a massive raw foundation. We expect the open community to quickly step in with specialized kernels, quantizations, and chat fine-tunes.