nanochat RLM d24 — supervised fine-tuned (SFT) checkpoint
Recurrent latent memory (RLM) model from the nanochat fork
(branch h100-rlm). The model reuses a 24-layer GPT backbone across bounded 1023-token text chunks and passes
512 learned memory vectors between chunks. Each cycle input is laid out as
[memory vectors] [<EOM>] [text tokens] [summary slots].
This repository holds the chat (SFT) checkpoint, the matching tokenizer, and the SFT run logs.
Model
| Setting | Value |
|---|---|
| Layers / heads / width | 24 / 12 / 1536 |
| Vocabulary | 32,768 (nanochat BPE, tokenizer.pkl) |
| Text chunk length | 1023 tokens |
| Memory vectors | 512 |
| Cycles per training episode | 4 (4,092 text tokens) |
| Weights | bfloat16 master weights, PyTorch state_dict |
Training
Pretraining: 5,574 steps of 1,047,552 tokens on ClimbMix (about 5.84B text tokens) on a single H100 NVL 94 GB.
SFT: one epoch over the nanochat chat mixture, loss on assistant tokens only. End-of-memory (<EOM>) readouts
are supervised only where the following chunk starts inside an assistant turn.
| Data | Conversations |
|---|---|
| SmolTalk (train) | 460,341 |
| MMLU auxiliary_train x3 | 299,526 |
| GSM8K train x4 | 29,892 |
| Total | 789,759 |
| SFT metric | Value |
|---|---|
| Updates | 504 (batch 1,047,552 tokens) |
| Held-out loss before SFT | 1.511 |
| Held-out loss after SFT | 0.867 |
| Wall clock | 14.5 h on one shared H100 NVL |
Learning rates were inherited from pretraining and scaled by 0.8, with no warmup, a 50% linear warmdown to 0,
and weight decay 0. Full configuration is in sft/run.json; per-update losses are in sft/updates.jsonl.
Files
tokenizer/tokenizer.pkl nanochat BPE tokenizer (sha256 387cfc08...)
tokenizer/token_bytes.pt
sft/checkpoints/model_000504.pt final SFT weights (state_dict)
sft/checkpoints/meta_000504.json model config, run contract, memory contract
sft/run.json SFT launch configuration
sft/result.json final metrics
sft/updates.jsonl per-update training log
The layout mirrors a nanochat training_run/<name>/ directory, so a downloaded snapshot can be used directly as a run.
Usage
git clone --branch h100-rlm https://github.com/ifsheldon/nanochat.git && cd nanochat
hf download rlm-exps/nanochat-rlm --local-dir training_run/rlm-sft
NANOCHAT_BASE_DIR=training_run/rlm-sft \
uv run --locked --python 3.10.21 --extra gpu --no-default-groups \
python -m scripts.rlm_sample --checkpoint-dir training_run/rlm-sft/sft/checkpoints \
--prompt "The capital of France is" --max-tokens 64
rlm_sample encodes the prompt as plain text, so chat-formatted prompting needs the special-token ids
(tokenizer.encode_special("<|user_start|>") and friends) assembled around the encoded user text.
In Python:
import torch
from pathlib import Path
from nanochat.rlm_inference import load_rlm_checkpoint # requires NANOCHAT_BASE_DIR set as above
model, tokenizer, meta = load_rlm_checkpoint(Path("training_run/rlm-sft/sft/checkpoints"), 504, torch.device("cuda"))
Note: the chat format uses nanochat special tokens (<|user_start|>, <|user_end|>, <|assistant_start|>,
<|assistant_end|>). This checkpoint is not a Hugging Face transformers model and will not load with AutoModel.