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

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

                    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

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:

TransformerConfig — the inner transformer:

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

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:

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

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

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 AssertionErrors.
• 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)

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

Key constants and defaults

Dependencies

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