The whole stack, in the open.
Everything that trains Apex-3, from the model to the data pipeline.
- Repository
- DW-dev-UE/LLM-from-scratch
- Path
- apex3/
- Commit
- 1d491e2
- Language
- Python · PyTorch
- Runs on
- 1× GH200 (ARM64, NGC PyTorch 26.09 container)
- Trains
- Apex-3
The complete pretraining stack behind Apex-3, released with the paper at a pinned commit.
It is the Python package apex3/: the mixture-of-experts model, the Muon optimizer, one-GPU offload into Grace memory, FP8 and a fused loss, DiLoCo for a second GH200, the data pipeline and an agent runtime.
Comments and docstrings are stripped. With docstrings excluded, every Python file parses to the same syntax tree as the server code behind the paper's step 10 benchmark and its training numbers after 10B tokens. Clone it below, or read it file by file in the explorer.
git clone https://github.com/DW-dev-UE/LLM-from-scratch.git cd LLM-from-scratch/apex3
What’s inside
Eight parts of the stack, each with the files that implement it.
| Component | Where | What it does |
|---|---|---|
| One-GPU offloadMemory | The FP32 masters and optimizer state live in pinned Grace memory, reached over NVLink-C2C. Chosen activations are stored on the CPU too. | |
| Fine-grained MoEModel | A sigmoid router picks the top experts, and a per-expert bias balances the load without a loss term. Experts run as grouped GEMMs. | |
| Muon + AdamWOptimizer | Muon on the hidden matrices, AdamW on the embeddings, LM head, router and one-dimensional parameters, with the WSD schedule. | |
| FlashAttention-3, varlenAttention | Document boundaries become varlen sequence lengths, so attention never crosses documents. The build script makes the FlashAttention-3 wheel. | |
| FP8 and a fused lossPrecision | FP8 through torchao, and a Triton kernel for the cross-entropy gradient that includes the z-loss term. | |
| DiLoCo for a second GH200Scale-out | Each worker publishes its difference from the shared weights, the differences are averaged, and an outer update with momentum follows. | |
| Data pipelineData | One module per stage: ingest, cleaning, decontamination, deduplication, redaction, repository packing and tokenization. | |
| Agent runtimeTools | Chat template, tools, a Docker sandbox, an MCP client, Skills, a KV-cache backend and the agent loop. |
Select a path to open it in the file explorer below.
Find your way around
Select a file to read what it does, then open it at the pinned commit.
Paths are relative to the repository root. Every link opens the pinned commit, 1d491e2.
Run it yourself
The container, the benchmark and the training command, converted from the paper’s appendix. Run them from the apex3/ directory.
# GH200 is ARM64: use the NGC PyTorch container. # Mount your code, data and checkpoint directories. The paper ran its commands with docker exec. docker run --gpus all --ipc=host --cap-add SYS_NICE \ --ulimit memlock=-1 --ulimit stack=67108864 \ -v CODE_DIR:CODE_DIR -v DATA_DIR:DATA_DIR -v CKPT_DIR:CKPT_DIR \ nvcr.io/nvidia/pytorch:26.09-py3
Image digest sha256:6e8ccc607fc2a51e3741667b86316a0889418ba8b78ed0ba96e8d282b648a7a6. FlashAttention-3 is built from source with server/build_fa3.sh, and the installation steps are in server/b_setup.sh. The default seccomp profile blocks set_mempolicy, so the container needs --cap-add SYS_NICE. Do not use --cpuset-mems 0: the GPU is no longer visible.
Linux exposes the GH200's HBM as NUMA node 1. Launch every process under numactl --membind=0 and drop the page cache before each run, or pinned host memory can spill into node 1 and end in GPU out-of-memory errors.
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True PYTHONPATH=.. \ TORCH_LOGS=recompiles,graph_breaks \ numactl --membind=0 torchrun --nproc_per_node 1 scripts/bench_mfu.py --config C \ --set n_experts=64 expert_ffn=1280 top_k=4 router_bias_rate=0.01 \ router_bias_update=proportional \ --offload --offload-accum cpu --compile-blocks --recompute none --act-budget 0.15 \ --moe-chunk 0 --micro-bs 6 --attn fa3 --accum 8 --ce-impl triton --ce-chunk 4096 \ --steps 4 --warmup 2 --gemm-offload hooks --act-offload 72 \ --precision fp8_rowwise --fp8-recipe tensorwise --json OUT.jsonl
The step 10 benchmark: the base flag line of server/bench_plan.sh plus the step 10 flags, with that line's --act-budget 0.05 changed to 0.15. It uses router_bias_rate=0.01, while training used 1e-3. It draws random tokens from a fixed seed and makes each row one document, so it needs no corpus.
The step 10 record shows micro_s_median 2.078 s, mfu_at_batch 0.5074, offloaded_tensors 118, offloaded_gib 56.4 and overflow_gib 0. This is a benchmark record, not a training figure. With TORCH_LOGS=recompiles,graph_breaks, it should show no graph breaks. PYTHONPATH=.. is for the public repository layout; the servers' benchmark used PYTHONPATH=$PWD.
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True PYTHONPATH=.. \ numactl --membind=0 torchrun --nproc_per_node 1 scripts/train.py \ --config C --eos-id 203776 --data NAME=WEIGHT:GLOB ... --val-data NAME=GLOB ... \ --set n_experts=64 expert_ffn=1280 top_k=4 router_bias_rate=1e-3 \ router_bias_update=proportional seq_aux_alpha=1e-3 \ --offload --offload-accum cpu --compile-blocks --recompute none --act-budget 0.15 \ --moe-chunk 0 --micro-bs 6 --attn fa3 --ce-impl triton --ce-chunk 4096 \ --gemm-offload hooks --act-offload 72 --precision fp8_rowwise --fp8-recipe tensorwise \ --total-tokens 1.1e12 --stop-at-tokens 0.51e12 --ckpt-dir CKPT --ckpt-every 1e9 \ --keep-ckpts 2 --arith-probe PROBE --arith-every 50 \ --diloco-workers 2 --diloco-rank R --diloco-dir DIR --diloco-peers PEER --diloco-key KEY
The command in effect after 10B tokens, launched after dropping the host page cache. The peer host, key and directories have been generalized. The launcher server/train_gh200.sh appends flags, so some appear more than once and argparse takes the last value. --seq-len 8192 and --seed 0 are defaults. PYTHONPATH=.. is for the public repository layout; on the servers, training ran from an installed package.
The Stage 1 mixture is code 0.35, web 0.40, math 0.06, papers 0.06, PDF 0.05, medical 0.03, textbooks 0.02, QA 0.02 and books 0.01. The 2.2 TB training corpus is not released. The arithmetic probe runs every 50 steps inside the training loop, so its time counts toward the logged wall-clock MFU. Logs showing CPU accumulation kernel not built or Triton CE unavailable mean a fallback path ran instead of the intended one.
Questions
What hardware do I need?
One GH200. It is ARM64, so use the NGC PyTorch 26.09 container shown above, with FlashAttention-3 built from server/build_fa3.sh. A second GH200 is optional: it joins through DiLoCo, as Apex-3’s second GPU worker did at 4B tokens. The step 10 benchmark needs no corpus.
Is the training data included?
No. The 2.2 TB training corpus is not released. The data pipeline code is: apex3/datapipe/ and apex3/scripts/data/pipeline.py cover download, ingest, decontam, clean, dedup, repo-pack, tokenize, long and stats. The Stage 1 mixture is listed under the Train tab above, and the benchmark draws random tokens from a fixed seed, so it runs without any data.
Where are the logs and the paper?
Every step of the run is on Training logs. The 20-page paper is on Research, and it also reports what did not work. The evidence bundles and figure scripts are in the repository under apex3/paper/.