Grok-1 User Guide

Welcome to the user guide for Grok-1, the open-weights release from xAI. This guide walks you through what this repository actually contains, how to obtain the model weights, and how to run the example code to generate text from the model on your own hardware.

Important

This repository is JAX example code for loading and running the Grok-1 open-weights model — it is not a chatbot, an API server, or a fine-tuning toolkit. It contains exactly four Python source files whose job is to load the released checkpoint and sample a single completion from a test prompt. The "Grok" product you may have used elsewhere is a separate, hosted service. This guide documents what the code in this repository does, nothing more.

1. Quick Start

If you already have a multi-GPU machine with enough memory and the weights downloaded, here is the entire workflow:

# 1. Clone the repository
git clone https://github.com/xai-org/grok-1.git && cd grok-1

# 2. Install the four Python dependencies
pip install -r requirements.txt

# 3. Place the downloaded ckpt-0 directory inside ./checkpoints/
#    (see "Downloading the Weights" below)

# 4. Run the example
python run.py

Expected output (the exact completion is non-deterministic across hardware, but the format is fixed):

Output for prompt: The answer to life the universe and everything is of course <generated continuation up to 100 tokens>

Warning

Grok-1 has 314 billion parameters. You cannot run this on a single consumer GPU, a laptop, or a CPU in any practical sense. A machine with enough aggregate GPU memory across multiple accelerators is required. See System Requirements before you begin.

2. What This Repository Is (and Is Not)

Grok-1 is an autoregressive, transformer-based large language model. This repository gives you the reference implementation needed to load the published weights and run a forward pass / sampling loop.

What you get:

• A faithful JAX + Haiku implementation of the Grok-1 architecture (model.py).
• A checkpoint loader that reconstructs the 314B-parameter state from sharded tensor files (checkpoint.py).
• An inference runner that builds the device mesh, loads weights, and exposes a token-sampling generator (runners.py).
• A runnable example that ties it all together and prints one completion (run.py).

What you do NOT get (so you don't go looking for it):

• Training or fine-tuning code — there is none. The code path is inference/sampling only.
• A command-line interface with flags — run.py takes no arguments. The prompt is hard-coded in the source.
• A REST API, web UI, or chat server — InferenceRunner.run() is an in-process Python generator, not a network service.
• A quantization tool — the code loads pre-quantized 8-bit weights; it does not quantize anything for you.

Note

The repository README states plainly: "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." In other words, this is a correctness reference, not a production-optimized serving engine. Slow generation is expected and by design.

Model specifications

These values are read directly from the model configuration in run.py and the repository README:

Note

Grok-1 uses grouped-query attention: 48 query heads share only 8 key/value heads. This is why the model can keep a 314B parameter count manageable for attention while supporting an 8,192-token context.

3. System Requirements

The single hard requirement is GPU memory — enough of it to hold a 314B-parameter model plus the key/value cache for an 8,192-token context.

Caution

The checkpoint loader copies tensor files through /dev/shm/ (Linux shared memory) during loading. On systems without /dev/shm (e.g. native Windows or a Linux box with a tiny tmpfs), loading will fail or stall. Run on Linux with a generously sized /dev/shm.

Tip

The example's device mesh is set to local_mesh_config=(1, 8), meaning it expects 8 local devices for model-parallel sharding. If your machine has a different number of accelerators, you will need to adjust this value — see Adapting the Example.

4. Installation

Step 1 — Get the code

git clone https://github.com/xai-org/grok-1.git
cd grok-1

Step 2 — Install dependencies

The repository pins exactly four dependencies in requirements.txt:

pip install -r requirements.txt

This installs:

Important

The JAX requirement uses a CUDA 12 wheel index: jax[cuda12-pip]==0.4.25 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html. Installing into an environment without a matching CUDA 12 toolchain will give you a CPU-only JAX, which cannot run the 314B model. Verify your CUDA setup first.

Tip

Use a clean virtual environment (python -m venv venv && source venv/bin/activate) or a fresh conda env. The pinned versions are specific; mixing them with an existing JAX install often causes version conflicts.

5. Downloading the Weights

The code ships without the model weights — the checkpoints/ directory is empty except for a placeholder README. You must download the ckpt-0 checkpoint and place it inside checkpoints/ so the loader can find it at ./checkpoints/ckpt-0.

There are two official ways to get the weights.

HuggingFace Hub

pip install huggingface_hub[hf_transfer]
huggingface-cli download xai-org/grok-1 \
    --repo-type model \
    --include ckpt-0/* \
    --local-dir checkpoints \
    --local-dir-use-symlinks False

This downloads the checkpoint directly into checkpoints/ckpt-0/.

Torrent (magnet link)

Use any torrent client with the magnet link published in the repository README:

magnet:?xt=urn:btih:5f96d43576e3d386c9ba65b883210a393b68210e&tr=https%3A%2F%2Facademictorrents.com%2Fannounce.php&tr=udp%3A%2F%2Ftracker.coppersurfer.tk%3A6969&tr=udp%3A%2F%2Ftracker.opentrackr.org%3A1337%2Fannounce

After the download completes, move the resulting ckpt-0 directory into checkpoints/.

After either method, your directory layout must look like this:

grok-1/
├── checkpoints/
│   └── ckpt-0/          ← the downloaded checkpoint shards live here
├── run.py
├── model.py
├── runners.py
├── checkpoint.py
└── tokenizer.model

Warning

The download is very large (the 8-bit checkpoint of a 314B model). Ensure you have hundreds of gigabytes of free disk before starting, and expect a long download even on fast connections.

Note

The SentencePiece tokenizer (tokenizer.model) is already included in the repository — you do not download it separately. The example loads it from ./tokenizer.model.

6. Running the Example

With dependencies installed and checkpoints/ckpt-0 in place, run:

python run.py

What happens, step by step

1. run.py builds the Grok-1 model configuration (all the specs from the table above are set here as Python values).
2. An InferenceRunner is created with the checkpoint path, the tokenizer path, and the device mesh layout.
3. inference_runner.initialize() builds the device mesh, compiles the forward and sampling functions, and loads the 314B-parameter checkpoint into device memory. This is the slow part.
4. The runner samples a completion for the built-in prompt and prints it.

The hard-coded test prompt in run.py is:

The answer to life the universe and everything is of course

and it is sampled with max_len=100 (generate up to 100 tokens) and temperature=0.01 (near-greedy, almost deterministic).

Note

Initialization logs progress at INFO level (the example calls logging.basicConfig(level=logging.INFO)). You will see messages like Loading checkpoint at ..., Initializing mesh ..., Compiling..., and Done compiling. before any output appears. Compilation and checkpoint loading dominate the wall-clock time; generation itself is comparatively brief.

Reading the output

The script prints a single line:

Output for prompt: The answer to life the universe and everything is of course <continuation>

Because temperature=0.01, the continuation is effectively greedy and will be nearly the same on each run for the same hardware.

7. Adapting the Example

Because there is no CLI, you customize behavior by editing run.py (and, for advanced changes, runners.py). This is the intended way to use the repository: it is example code meant to be read and modified.

Change the prompt

In run.py, edit the inp variable inside main():

inp = "Write a haiku about open-source models:"
print(f"Output for prompt: {inp}", sample_from_model(gen, inp, max_len=100, temperature=0.01))

Change generation length and temperature

The sample_from_model() call accepts max_len and temperature:

# Longer, more creative output
sample_from_model(gen, inp, max_len=256, temperature=0.7)

Note

Under the hood, a sampling request also has a nucleus_p (top-p) value. The convenience function sample_from_model() hard-codes nucleus_p=1.0 (no nucleus filtering) and rng_seed=42. To control top-p sampling or the random seed, construct a Request directly (see the Request dataclass in runners.py) and call the generator's .send() method yourself, mirroring what sample_from_model() does.

Match the mesh to your hardware

The example assumes 8 local devices. The relevant line in run.py is:

local_mesh_config=(1, 8),   # (data_parallel, model_parallel)

The second number is the model-parallel degree and should match your accelerator count. Adjust it together with bs_per_device (set to 0.125 in the example) so the global batch size remains valid for your setup.

Caution

The model is sharded across the mesh using fixed partition rules. Choosing a mesh shape that does not evenly divide the model dimensions will produce sharding errors at compile time. The provided (1, 8) configuration is the tested layout.

8. Troubleshooting

The example surfaces a small number of explicit errors. Below are the ones you are most likely to hit, with their exact messages.

Error: Parameters in the code are not matching checkpoint parameters.

Parameters in the code are not matching checkpoint parameters.
Params missing in checkpoint: {...}
Params missing in code: {...}

• Cause: The checkpoint you placed in checkpoints/ckpt-0 does not structurally match the model defined in run.py — typically a wrong, partial, or corrupted download, or a modified model configuration.
• Solution: Re-download the official ckpt-0 and confirm it sits at exactly ./checkpoints/ckpt-0. Do not change the model architecture values in run.py unless you also have a matching checkpoint.
• Prevention: Use the HuggingFace --include ckpt-0/* command exactly as shown so the directory structure is preserved.

Error: Input must not be scalar.

ValueError: Input must not be scalar.

• Cause: A tensor with no shape was passed into a linear or routing layer — almost always a sign of an incorrectly shaped prompt or a manual modification to the input pipeline.
• Solution: Ensure prompts are non-empty strings. If you wrote custom code that constructs token arrays, confirm they are at least 1-D.
• Prevention: Use sample_from_model() or a properly constructed Request; let the runner handle tokenization.

Error: Mask dimensionality ... must match logits dimensionality ...

ValueError: Mask dimensionality {n} must match logits dimensionality {m} ...

• Cause: An attention mask with the wrong number of dimensions reached the attention computation, generally from custom changes to masking or memory handling.
• Solution: Revert custom changes to the attention/memory code and reproduce with the stock example first.
• Prevention: Treat model.py's attention path as a reference; if you must change it, keep mask shapes at 4 dimensions as the assertions require.

Out-of-memory during checkpoint load or compilation

• Cause: Insufficient aggregate GPU memory for a 314B-parameter model.
• Solution: Run on a machine with more/larger accelerators. There is no built-in CPU-offload or low-memory mode in this repository.
• Prevention: Confirm your hardware can hold the full checkpoint before downloading hundreds of GB of weights.

Loading hangs or fails referencing /dev/shm

• Cause: The checkpoint loader stages tensors through Linux shared memory (/dev/shm). A missing or too-small /dev/shm breaks this.
• Solution: Run on Linux and enlarge /dev/shm (e.g. mount tmpfs with sufficient size).
• Prevention: Use a standard Linux GPU host; avoid native Windows for this workload.

JAX runs on CPU instead of GPU

• Cause: JAX was installed without the CUDA 12 wheel, or the CUDA runtime is missing.
• Solution: Reinstall using the exact requirements.txt (which points at the JAX CUDA release index) inside an environment with a working CUDA 12 toolchain.
• Prevention: Verify jax.devices() lists GPU devices before attempting a run.

9. Frequently Asked Questions

Q: Can I run Grok-1 on my single RTX 4090 / a laptop / a CPU? No. The model has 314B parameters and requires a multi-accelerator machine with enough collective GPU memory. There is no quantization-down or offload path in this repository.

Q: Is this the same as the Grok chatbot? No. This repository is the open-weights model plus minimal example code to run inference. The hosted Grok assistant is a separate product. This code does not include a chat interface.

Q: How do I serve this as an API? This repository does not provide a server. InferenceRunner.run() returns an in-process generator you drive with Request objects from Python. Building a network service around it is left to you.

Q: Can I fine-tune or continue training Grok-1 with this code? No. The repository contains inference/sampling code only. There is no training loop, optimizer, or dataset pipeline.

Q: Why is generation so slow? By design. The README states the MoE layer implementation is intentionally unoptimized to avoid custom kernels and to keep the code a clear correctness reference. It is not built for throughput.

Q: Where is the prompt configured? In run.py, as the inp variable inside main(). There are no command-line flags — edit the source to change it.

Q: What context length is supported? Up to 8,192 tokens total (prompt plus generated output combined).

Q: Do I need to download the tokenizer separately? No. tokenizer.model (a SentencePiece model) ships in the repository and is loaded from ./tokenizer.model.

Q: What license covers the code and weights? Both the source files in this repository and the Grok-1 model weights are released under the Apache 2.0 license.

Q: How do I change the random seed or use top-p (nucleus) sampling? Construct a Request directly with your desired rng_seed and nucleus_p, then send it into the runner's generator, instead of using the sample_from_model() convenience wrapper (which fixes rng_seed=42 and nucleus_p=1.0).

10. Reference

Files in this repository

Default example parameters

Tip

For the architecture deep-dive — how the MoE routing, grouped-query attention, sharding, and checkpoint loading actually work — see the companion Grok-1 Admin Guide.