# 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:

```shell
# 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):

```text
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](#3-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](https://github.com/google-deepmind/dm-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:

| Specification | Value |
|---|---|
| Parameters | 314B |
| Architecture | Mixture of Experts (MoE) |
| Number of experts | 8 |
| Experts used per token | 2 |
| Layers | 64 |
| Query attention heads | 48 |
| Key/value attention heads | 8 |
| Embedding size | 6,144 (`48 × 128`) |
| Vocabulary size | 131,072 (`128 × 1024`) |
| Maximum context length | 8,192 tokens |
| Positional encoding | Rotary embeddings (RoPE) |
| Weight format | 8-bit quantized (`QuantizedWeight8bit`) |
| Tokenizer | SentencePiece (`tokenizer.model`, shipped in the repo) |

> [!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.

| Requirement | Detail |
|---|---|
| Accelerators | Multiple GPUs (the example is configured for an 8-way model-parallel mesh) |
| GPU memory | Must collectively hold a 314B-parameter 8-bit checkpoint plus activations and KV cache |
| CUDA | CUDA 12.x (the pinned JAX wheel is `jax[cuda12-pip]`) |
| Python | A recent CPython (3.10+ recommended for the JAX/Haiku versions pinned) |
| Disk | Hundreds of GB free for the downloaded weights (`ckpt-0`) |
| OS | Linux is strongly recommended — the checkpoint loader uses `/dev/shm` (shared memory) for fast tensor loading |

> [!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](#6-adapting-the-example).

---

## 4. Installation

### Step 1 — Get the code

```shell
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`:

```shell
pip install -r requirements.txt
```

This installs:

| Package | Pinned version | Role |
|---|---|---|
| `dm_haiku` | `0.0.12` | Neural-network module library (the `hk.Module` layers) |
| `jax[cuda12-pip]` | `0.4.25` | Array/autodiff backend with CUDA 12 GPU support |
| `numpy` | `1.26.4` | Host-side array handling and checkpoint tensors |
| `sentencepiece` | `0.2.0` | Tokenizer runtime (encodes prompts, decodes output) |

> [!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"

    ```shell
    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:

    ```text
    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:

```text
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:

```shell
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:

```text
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:

```text
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()`:

```python
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`:

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

| Parameter | Meaning | Notes |
|---|---|---|
| `max_len` | Maximum number of tokens to generate | Total context (prompt + output) cannot exceed 8,192 |
| `temperature` | Sampling randomness | `0.01` ≈ greedy/deterministic; higher = more varied |

> [!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:

```python
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.`

```text
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.`

```text
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 ...`

```text
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

| File | Purpose |
|---|---|
| `run.py` | Entry point — builds the model config and runs the example |
| `model.py` | The Grok-1 architecture (attention, MoE, transformer, embeddings) |
| `runners.py` | Inference runner, device mesh setup, sampling logic |
| `checkpoint.py` | Loads the sharded 314B checkpoint from disk |
| `tokenizer.model` | SentencePiece tokenizer (shipped) |
| `requirements.txt` | The four pinned dependencies |
| `checkpoints/` | Where you place the downloaded `ckpt-0` weights |

### Default example parameters

| Setting | Value (in `run.py`) |
|---|---|
| Checkpoint path | `./checkpoints/` |
| Tokenizer path | `./tokenizer.model` |
| Device mesh (local) | `(1, 8)` |
| Batch size per device | `0.125` |
| Pad bucket size | `1024` |
| Test prompt | `"The answer to life the universe and everything is of course"` |
| `max_len` | `100` |
| `temperature` | `0.01` |

> [!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**.

---

_Generated June 2026 from xai-org/grok-1 at 7050ed2. AI-generated documentation. The library is checked by sampling (three repositories so far); this guide was not individually reviewed. Verify against the project before relying on it. Unofficial; not affiliated with or reviewed by the project._