The Problem

Running a 314B-parameter Mixture-of-Experts model is not a pip-install-and-go exercise. Grok-1 needs multi-GPU memory far beyond a single device, a checkpoint that ships via torrent or HuggingFace, and JAX sharding logic to distribute the weights. This repo exists to prove the open-weights release actually runs, not to be a production serving stack.

What This Does

This is the official JAX example code for loading and sampling from the Grok-1 open-weights model. model.py defines the full architecture — 64 layers, 8 experts with 2 active per token, RoPE embeddings, activation sharding, and 8-bit quantization support. checkpoint.py handles loading the large checkpoint, including shared-memory and fast-pickle helpers. runners.py contains the sampling and inference logic (top_p_filter, sample_token, make_forward_fn), and run.py is the single entry point.

The README is explicit that the MoE implementation is deliberately inefficient — it avoids custom kernels to validate correctness without a specialized build.

How It Is Wired

Execution starts at main in run.py:24, which reaches 22 functions. It calls run in runners.py:442, which drives make_forward_fn, prefill_memory, and pad_to_size — this is where the model actually executes against the loaded checkpoint.

model.py is the hub: 3 modules import it, it holds 55 functions and 22 classes, and it is 1,145 lines. The hottest internal edges confirm its centrality: __call__ calls with_sharding_constraint 15 times and layer_norm 6 times. KVMemory and Memory are each called from 4 distinct places, so changes to those ripple widest. runners.py has instability 0.67 (imports more than it is imported by), and run.py has instability 1.0 — a pure leaf that nothing else calls.

The only external effect is file I/O: checkpoint.py reads and writes files (loading tensors, copying to shared memory). There is no database or network call in the traced paths.

File-by-file map: model.py owns the architecture and sharding; runners.py owns inference and sampling; checkpoint.py owns weight loading; run.py owns the orchestration.

How To Use It

Setup: pip install -r requirements.txt (the only declared dependency manager is pip; pyproject.toml exists but has no lockfile).

Configuration: None — no environment variables or config files. The tokenizer is tokenizer.model in the repo root.

Running it:

pip install -r requirements.txt
python run.py

Before running, download the checkpoint and place the ckpt-0 directory in checkpoints/ — via the magnet link in the README or huggingface-cli download xai-org/grok-1 --repo-type model --include ckpt-0/* --local-dir checkpoints --local-dir-use-symlinks False. A machine with enough GPU memory for 314B parameters is required.

Real-World Use

This fits as a verification harness: after downloading weights, run python run.py to confirm the checkpoint loads and the model produces a coherent sample on a test input. It is not a serving path — no batching, no KV-cache optimization, no inference server. Its value is confirming a checkpoint is intact and the architecture is correctly wired before you invest in a custom-kernel implementation.

Code Health & Issues

Static analysis (measured, not opinion) found 4 medium findings across 3 kinds:

  • Med — deep nesting — model.py, runners.py: max indentation depth 6; control flow is hard to follow. Fix with early returns and extracted blocks.
  • Med — oversized file — model.py: 1,145 code lines; changes ripple widely. Split by responsibility.
  • Med — duplicated code blocks — checkpoint.py, model.py, run.py, runners.py: 5 repeated 6-line blocks across 4 files. Extract shared helpers.

The code-health audit adds one high finding: no lockfile beside pyproject.toml, so the artifact you test can differ from what you ship — commit the generated lockfile. Medium: no Dependabot/Renovate configured for the 2 manifests.

SDLC gaps visible in the structure: no test files, no CI/CD pipeline, no Dockerfile. License is present (Apache 2.0), no committed secrets detected.

The Bottom Line

This is a faithful, minimal reference implementation of Grok-1 — correct by design, slow by choice, and honest about it. Use it to validate weights and understand the architecture, not to serve traffic. The code quality is typical of research releases: a large monolithic model.py, no tests, no lockfile. If you plan to build on it, budget for refactoring and dependency pinning first.