# Grok-1 Admin & Architecture Guide

This guide is for engineers, ML architects, and operators who need to understand how the Grok-1 reference implementation is built — the model architecture, the parallelism strategy, the checkpoint format, and the operational characteristics of running a 314B-parameter Mixture-of-Experts model. It is written to be exhaustive enough that a competent engineer could reconstruct or extend the implementation from this document alone.

---

## 1. Executive Summary

`xai-org/grok-1` is the **open-weights release** of xAI's Grok-1 large language model, accompanied by a compact JAX + [Haiku](https://github.com/google-deepmind/dm-haiku) reference implementation. The repository's purpose is narrow and deliberate: load the released 314B-parameter checkpoint and run a forward/sampling pass to validate model correctness.

| Attribute | Value |
|---|---|
| Model family | Grok-1 (autoregressive transformer LM) |
| Total parameters | 314B |
| Architecture | Mixture of Experts (8 experts, 2 active per token) |
| Framework | JAX `0.4.25` + dm-haiku `0.0.12` |
| Source files | 4 Python modules (`model.py`, `runners.py`, `checkpoint.py`, `run.py`) |
| Total source size | ~77 KB |
| Parallelism | SPMD sharding via JAX `Mesh` + `shard_map`/`pjit` |
| Weight format | 8-bit quantized (`QuantizedWeight8bit`) |
| License | Apache 2.0 (code **and** weights) |
| Source commit | `7050ed2` (2024-03-19) |

> [!IMPORTANT]
> This is **inference reference code**, not a serving stack. There is no training loop, no fine-tuning, no REST/gRPC server, no CLI, and no autoscaling. The MoE implementation is explicitly unoptimized — the README states it "was chosen to avoid the need for custom kernels to validate the correctness of the model." Treat throughput as a non-goal of this repository. Any production deployment requires building substantial infrastructure around this core.

> [!NOTE]
> Because the user-facing surface is a single hard-coded example (`run.py`), most of the engineering value here is in the *architecture the code reveals*. This guide therefore weights heavily toward the model internals, sharding, and checkpoint subsystems.

---

## 2. Architecture Overview

### High-level data flow

```text
                    run.py  (main)
                       │  builds LanguageModelConfig(TransformerConfig(...))
                       ▼
              InferenceRunner.initialize()
                       │  ┌───────────────────────────────────────────┐
                       │  │ make_mesh((1,8),(1,1)) → jax Mesh           │
                       │  │ hk.transform(forward / sample / prefill)    │
                       │  │ jax.eval_shape → partition rules → sharding │
                       │  │ checkpoint.restore() → 314B params on mesh  │
                       │  └───────────────────────────────────────────┘
                       ▼
              InferenceRunner.run()  (Python generator)
                       │  yields ready → receives Request(prompt,...)
                       ▼
   tokenizer.encode → prefill_memory (pjit) → sample_step (pjit) loop
                       │  per-step: LanguageModel.__call__ → sample_token
                       ▼
              tokenizer.decode → output string
```

### Component inventory

| Module | Responsibility | Key symbols |
|---|---|---|
| `model.py` | The neural network: attention, MoE, transformer stack, embeddings, quantized weights, partition rules | `LanguageModel`, `Transformer`, `DecoderLayer`, `MultiHeadAttention`, `MoELayer`, `Router`, `DenseBlock`, `RMSNorm`, `RotaryEmbedding`, `InOutEmbed`, `QuantizedWeight8bit`, `TransformerConfig`, `LanguageModelConfig` |
| `runners.py` | Orchestration: device mesh, JIT compilation, KV-cache memory, the sampling loop | `InferenceRunner`, `ModelRunner`, `sample_token`, `top_p_filter`, `make_mesh`, `sample_from_model`, `Request`, `SampleSettings`, `SampleOutput` |
| `checkpoint.py` | Loading sharded tensors from disk into the distributed state | `restore`, `load_tensors`, `replace_with_load_state`, `fast_unpickle`, `copy_to_shm` |
| `run.py` | Entry point that wires the model config to the runner and prints one completion | `main` |

> [!NOTE]
> The dependency graph is acyclic and shallow: `run.py → runners.py → {model.py, checkpoint.py}`, and `checkpoint.py → model.py` (for `QuantizedWeight8bit`). All 27 classes and 80 functions extract cleanly because the codebase is idiomatic Python/Haiku.

---

## 3. The Model (`model.py`)

### Configuration objects

The model is described declaratively by two dataclasses, instantiated in `run.py`:

`LanguageModelConfig` — the outer wrapper:

| Field | Example value | Meaning |
|---|---|---|
| `vocab_size` | `131072` (`128 × 1024`) | Token vocabulary size |
| `pad_token` | `0` | Padding token id |
| `eos_token` | `2` | End-of-sequence token id |
| `sequence_len` | `8192` | Maximum context length |
| `embedding_init_scale` | `1.0` | Embedding init scaling |
| `embedding_multiplier_scale` | `78.38367...` | Multiplier applied to input embeddings |
| `output_multiplier_scale` | `0.57735...` | Multiplier applied to output logits |
| `fprop_dtype` | `jnp.bfloat16` | Forward-pass compute dtype |
| `shard_embeddings` | `True` | Whether to shard the embedding matrix |

`TransformerConfig` — the inner transformer:

| Field | Example value | Meaning |
|---|---|---|
| `emb_size` | `6144` (`48 × 128`) | Model hidden dimension |
| `widening_factor` | `8` | FFN expansion factor |
| `key_size` | `128` | Per-head key/query dimension |
| `num_q_heads` | `48` | Query heads (grouped-query attention) |
| `num_kv_heads` | `8` | Key/value heads |
| `num_layers` | `64` | Decoder layers |
| `attn_output_multiplier` | `0.08838...` | Attention logit scaling |
| `num_experts` | `8` | MoE experts per layer |
| `num_selected_experts` | `2` | Experts evaluated per token |
| `shard_activations` | `True` | Whether activations are model-sharded |
| `data_axis` / `model_axis` | `"data"` / `"model"` | Mesh axis names for sharding |

> [!TIP]
> These dataclass fields **are** the configuration reference for the model — there are no environment variables or config files. To change model behavior you edit these kwargs in `run.py`. The values above are the released Grok-1 configuration; changing them requires a matching checkpoint.

### Attention: grouped-query + RoPE

`MultiHeadAttention` implements grouped-query attention (GQA): 48 query heads are reshaped to share 8 key/value heads (`query_heads.reshape(b, t, kv_h, h // kv_h, d)`). Salient details:

- **Rotary embeddings** (`RotaryEmbedding`, RoPE, base exponent `10000`) are applied to both query and key heads, offset by the current KV-cache step for correct positions during incremental decoding.
- **Logit soft-capping:** attention logits are scaled by `attn_output_multiplier` and then passed through `max_attn_val * tanh(logits / max_attn_val)` with `max_attn_val = 30.0`, bounding them before softmax.
- **Softmax in fp32:** attention weights are always computed in `float32` for numerical stability, then cast back to the forward dtype.
- **Causal masking:** a lower-triangular mask is combined with the padding mask; masked positions are set to `-1e30` before softmax.

### Mixture of Experts (`MoELayer` + `Router`)

Each decoder layer with `num_experts > 1` routes tokens through a sparse MoE block:

1. `Router._compute_routing_prob` computes routing logits in fp32 and applies softmax → per-token expert probabilities.
2. `jax.lax.top_k` selects the top `num_selected_experts` (2) experts per token.
3. The expert FFN (`DenseBlock`, a gated GeLU MLP: `gelu(linear(x)) * linear_v(x) → linear_1`) is `vmap`-ped across all 8 experts via `hk.experimental.transparent_lift`.
4. Two `shard_map`-wrapped matmuls (`moe_slow_matmul1`/`moe_slow_matmul2`) compute expert outputs and combine them with one-hot selection masks, scaled by the routing gate.

> [!WARNING]
> The functions are literally named `moe_slow_matmul1` and `moe_slow_matmul2`. As the README states, this MoE path is **intentionally inefficient** — it computes dense expert matmuls and selects with one-hot masks rather than gathering, to avoid custom kernels. This is correct but compute-wasteful; it is the single biggest performance gap versus a production MoE kernel.

> [!NOTE]
> The FFN hidden size is computed by `ffn_size(emb_size, widening_factor)` as `int(widening_factor * emb_size) * 2 // 3`, rounded up to a multiple of 8 — the standard "2/3 widening" used for gated (SwiGLU-style) MLPs.

### Transformer stack and language model head

- `DecoderLayer` applies pre-/post-norm RMSNorm around the attention and FFN/MoE blocks with residual connections (`h += h_attn`, `h += h_dense`), each wrapped in sharding constraints.
- `Transformer.__call__` loops over all 64 layers, threading the per-layer `KVMemory` through for incremental decoding, and returns the final embeddings plus the updated `Memory`.
- `LanguageModel` embeds tokens via `InOutEmbed` (with **tied input/output weights** — `in_out_embed.decode()` reuses the embedding matrix transposed), applies the `embedding_multiplier_scale`, runs the transformer, applies a final RMSNorm, decodes to logits, and applies `output_multiplier_scale`.

### Quantized weights

`QuantizedWeight8bit` is a small dataclass holding `weight` (int8) and `scales` (the per-block dequant scales), registered as a JAX pytree node so it flows through `tree_map`/`tree_flatten`. The `Linear` and MoE matmul paths detect `hasattr(w, "scales")` and dequantize on the fly (`w.weight * w.scales`) inside a `shard_map`. This is how the 314B model fits in 8-bit storage while computing in bfloat16.

---

## 4. Inference Runtime (`runners.py`)

### `ModelRunner` — model lifecycle

`ModelRunner` owns model initialization and weight loading:

- `initialize()` sets `fprop_dtype = bfloat16`, computes batch sizes from `bs_per_device × local_gpus × replicas`, builds the mesh via `make_mesh`, and constructs the `hk.transform`-ed forward functions.
- `get_state_sharding()` uses `jax.eval_shape` + the model's `partition_rules()` (via `apply_rules`) to derive a `PartitionSpec` tree for every parameter **without materializing weights** — this is what makes loading a 314B model tractable.
- `load_or_init()` either initializes fresh parameters or calls `checkpoint.restore()` to load the checkpoint shaped to the derived sharding.

### `InferenceRunner` — compilation and sampling

`InferenceRunner.initialize()` does the heavy lifting:

- Loads the SentencePiece tokenizer from `tokenizer_path`.
- `hk.transform`s four functions: `hk_forward`, `hk_sample_step`, `hk_new_memory`, `hk_prefill_memory`.
- Computes parameter sharding via `eval_shape` + `apply_rules`.
- Wraps the step functions in `pjit.pjit` with explicit `in_shardings`/`out_shardings` and `donate_argnums` for the KV-cache memory (donation enables in-place buffer reuse).

`InferenceRunner.run()` is a **Python generator** implementing continuous batching:

1. Allocates a fixed-size batch of KV-cache memory and a free-slot list.
2. On each `yield`, accepts a `Request`, tokenizes the prompt, pads it to a bucket size (`get_pad_bucket`), and runs `prefill_memory` into a free batch slot.
3. Runs `sample_step` across the whole batch, harvests finished sequences (when `len(all_tokens) >= max_len`), decodes them, frees their slots, and yields the output string.

> [!NOTE]
> Padding uses fixed buckets (`pad_sizes=(1024,)`) so the compiled `prefill_memory` is reused across prompts of varying length. Prompts longer than the bucket are **left-truncated** (`pad_to_size` keeps the last `size` tokens). Be aware that very long prompts silently lose their beginning.

### Sampling (`sample_token` + `top_p_filter`)

`sample_token` divides logits by temperature, masks disallowed tokens to `-1e10`, applies nucleus (top-p) filtering, then draws with `jax.random.categorical`. It also returns the chosen token's probability and the top-8 candidate tokens/probabilities (`TOP_K = 8`) for inspection.

`top_p_filter` sorts logits, computes the cumulative softmax mass, and masks out tokens below the `1 - top_p` threshold (setting them to `-1e10`).

### Device mesh (`make_mesh`)

`make_mesh` builds a 2-D hybrid device mesh with axes `("data", "model")` using `mesh_utils.create_hybrid_device_mesh`, combining a `local_mesh_config` (intra-host) and `between_hosts_config` (inter-host). The example uses `(1, 8)` local — pure 8-way model parallelism on a single host.

---

## 5. Checkpoint Subsystem (`checkpoint.py`)

The checkpoint is a directory `ckpt-0/` of per-tensor pickle files named `tensor{NNNNN}_{NNN}` (tensor index + shard index). Loading is parallel and shared-memory-staged.

| Function | Role |
|---|---|
| `restore()` | Top-level entry: flattens the expected state shapes, loads tensors, unflattens, validates keys, converts host-local arrays to a global sharded array |
| `load_tensors()` | Loads tensors in parallel via a 32-worker `ThreadPoolExecutor`; only the shards owned by this process are read, others are zero-filled |
| `replace_with_load_state()` | Merges loaded tensors into an init-state pytree, applying optional rename/exclude regex rules per parameter path |
| `fast_unpickle()` / `fast_pickle()` | Pickle load/save that stages the file through `/dev/shm` |
| `copy_to_shm()` / `copy_from_shm()` | Context managers that copy a file into Linux shared memory for fast random access |

> [!CAUTION]
> Loading uses Python `pickle` (`pickle.load`) on checkpoint files, and `checkpoint.py` registers `QuantizedWeight8bit` into `sys.modules['__main__']` so the pickle can resolve it. **Unpickling executes arbitrary code embedded in the file.** Only load checkpoints from trusted sources (the official xAI release). A malicious `ckpt-0` could run code on your host. This is the single most important security consideration when operating this code.

> [!IMPORTANT]
> `restore()` performs a key-set sanity check between the checkpoint and the code-defined parameters and raises a descriptive `ValueError` listing params missing in the checkpoint vs. missing in code. This is your first line of defense against a partial/corrupted download or an architecture mismatch.

---

## 6. Sharding & Parallelism Deep Dive

Grok-1 is too large for one device; the entire design is built around JAX SPMD sharding across the `("data", "model")` mesh.

- **Partition rules** (`TRANSFORMER_PARTITION_RULES`, `LM_PARTITION_RULES`) are ordered `(regex-path → PartitionSpec)` pairs. `apply_rules` walks each parameter's pytree path, matches it against the rules, and assigns a `PartitionSpec`. Attention/FFN weights are sharded across both `data` and `model` axes; norms, biases, and the router are mostly replicated (`P(None)`).
- **Activation sharding** is applied throughout via `with_sharding_constraint` (a no-op when no physical mesh is present, enabling single-device debugging).
- **MoE sharding** uses `shard_map` with explicit in/out specs so the slow matmuls run distributed across the `model` axis with a `psum` reduction.
- **KV cache** (`Memory`/`KVMemory`) is sharded with `get_memory_sharding()` and donated to `pjit` so buffers are reused in place across decode steps.

> [!NOTE]
> Because sharding is rule-driven by parameter path regexes, renaming layers or changing the module hierarchy will silently drop parameters from the rules (they fall to `P(None)` / replicated) — watch the `"no matching found!"` log lines emitted by `apply_rules` when modifying the model.

---

## 7. Operations

### Running

There is one operational command:

```shell
python run.py
```

All behavior (model config, mesh, prompt, sampling) is set in `run.py`. Operational "configuration" means editing that file. Logging is enabled at `INFO` via `logging.basicConfig(level=logging.INFO)`; there is also a separate `"rank"` logger used for per-process diagnostics during loading and sharding.

### What to watch in the logs

| Log message | Phase |
|---|---|
| `Initializing mesh for ...` | Mesh construction |
| `Loading checkpoint at .../ckpt-0` | Checkpoint load start |
| `Restored from ckpt: ...` / `Not found in ckpt: ...` | Per-parameter load status |
| `Precompile {size}` | Prefill compilation per pad bucket |
| `Compiling...` / `Done compiling.` | Sample-step XLA compilation |

### Operational characteristics

- **Startup is dominated** by checkpoint loading and XLA compilation, not generation.
- **No persistence/state** beyond the loaded weights — each process run reloads everything.
- **No graceful shutdown / health endpoints** — it is a script.
- **Determinism:** seeds are fixed (`rng_seed=42` in the convenience path), but cross-hardware/XLA differences mean outputs are not bit-identical across machines.

> [!TIP]
> To stand up anything resembling a service, you would keep one `InferenceRunner.run()` generator alive (paying the load+compile cost once) and feed it `Request` objects from your own server loop. The generator already supports continuous batching across the fixed batch size — that is the natural integration seam.

---

## 8. Security Analysis

| Area | Finding | Recommendation |
|---|---|---|
| Checkpoint deserialization | `pickle.load` on checkpoint files; arbitrary code execution risk | Load only the official, integrity-verified checkpoint; never accept third-party `ckpt-0` directories |
| `__main__` injection | `sys.modules['__main__'].QuantizedWeight8bit` is set to enable unpickling | Acceptable for trusted checkpoints; do not generalize this pattern to untrusted inputs |
| Shared memory usage | Tensors staged through `/dev/shm` | Ensure `/dev/shm` is on a trusted, single-tenant host; multi-tenant nodes could expose staged tensors |
| Network exposure | None — no server, no listening sockets | If you build a server around it, add authentication, input limits, and rate limiting yourself |
| Input validation | Prompts are tokenized as-is; long prompts are silently left-truncated | Enforce prompt length limits at your integration layer |
| Secrets | None present; no credentials, keys, or tokens in the repository | N/A |

> [!CAUTION]
> The dominant security risk is **untrusted checkpoint loading via pickle**. There is no signature verification or safe-tensors path in this repository. If you cannot fully trust a checkpoint's provenance, do not load it with this code.

---

## 9. Code Quality & Conventions

- **Style:** formatted/linted with Ruff (`pyproject.toml`: line length 100, 4-space indent; a handful of lint rules ignored, e.g. `E731`/`F403`).
- **Idioms:** heavy use of Haiku modules (`hk.Module`, `hk.transform`, `hk.transparent`), JAX functional transforms (`vmap`, `pjit`, `shard_map`, `eval_shape`), and `NamedTuple`/`@dataclass` for typed containers.
- **Assertions as contracts:** the attention path uses dense `assert` statements on tensor shapes/dtypes — excellent as executable documentation, but they run unconditionally (no `-O` guard) and raise bare `AssertionError`s.
- **Numerical care:** routing and softmax are forced to fp32; logit soft-capping bounds attention; compute dtype is bfloat16. These reflect production-grade attention to stability.
- **Clarity over speed:** the codebase consistently chooses readable, kernel-free implementations (the MoE matmuls being the prime example), consistent with its stated role as a correctness reference.

### Extracted surface (for reference)

| Metric | Count |
|---|---|
| Source files analyzed | 4 |
| Classes | 27 |
| Functions | 80 |
| Explicit error raises | 4 (`ValueError`) |
| CLI commands | 0 (no CLI exists) |
| API endpoints | 0 (no server exists) |
| Config keys (env/file) | 0 (config is Python kwargs in `run.py`) |

---

## 10. Potential Improvements

These are observations from the source, framed as fork/extension opportunities — not defects in a reference implementation.

- [ ] **Optimize the MoE path.** Replace `moe_slow_matmul1/2`'s dense-then-one-hot approach with a gather/scatter or a fused expert kernel. This is the largest throughput lever and is explicitly acknowledged in the README.
- [ ] **Add a safe checkpoint format.** Migrate from `pickle` to `safetensors` (or add signature verification) to remove the arbitrary-code-execution risk during loading.
- [ ] **Parameterize `run.py`.** Add an `argparse`/config layer so the prompt, mesh shape, `max_len`, `temperature`, and `nucleus_p` can be set without editing source — currently the only way to change them.
- [ ] **Expose a serving seam.** Wrap the existing continuous-batching generator in a thin server (with auth, length limits, and backpressure) so the load+compile cost is amortized.
- [ ] **Guard prompt truncation.** Left-truncation of over-long prompts is silent; surface a warning or reject prompts exceeding the pad bucket / context length.
- [ ] **Make `/dev/shm` staging optional.** Provide a fallback path for environments without large shared memory to broaden portability beyond Linux GPU hosts.

---

## 11. Quick Reference

### Module map

| File | Lines of responsibility |
|---|---|
| `run.py` | Config + example entry point |
| `model.py` | Architecture (attention, MoE, transformer, embeddings, sharding rules, quantized weights) |
| `runners.py` | Mesh, compilation, KV-cache, sampling loop |
| `checkpoint.py` | Parallel sharded checkpoint loading via `/dev/shm` + pickle |

### Key constants and defaults

| Name | Value | Location |
|---|---|---|
| `CKPT_PATH` | `./checkpoints/` | `run.py` |
| Tokenizer path | `./tokenizer.model` | `run.py` |
| `local_mesh_config` | `(1, 8)` | `run.py` |
| `between_hosts_config` | `(1, 1)` | `run.py` |
| `bs_per_device` | `0.125` | `run.py` |
| `pad_sizes` | `(1024,)` | `run.py` |
| `TOP_K` | `8` | `runners.py` / `model.py` |
| `max_attn_val` | `30.0` | `model.py` (attention soft-cap) |
| RoPE base exponent | `10000` | `model.py` |
| Default rng seed | `42` | `runners.py` |

### Dependencies

| Package | Version |
|---|---|
| `dm_haiku` | `0.0.12` |
| `jax[cuda12-pip]` | `0.4.25` |
| `numpy` | `1.26.4` |
| `sentencepiece` | `0.2.0` |

### Operational facts

- **Entry point:** `python run.py` (no arguments).
- **Hardware:** multi-GPU, CUDA 12, Linux with sizable `/dev/shm`.
- **Weights:** download `ckpt-0` (HuggingFace `xai-org/grok-1` or magnet torrent) into `checkpoints/`.
- **License:** Apache 2.0 (code and weights).

---

_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._