Pre-training, step by step.

What the loss is

Line 4 of the loop produces one number, the loss. The rest of this page is about pushing that number down: faster, on more data, and without the run crashing. Here is what the number measures, one line at a time.

Line 3 gives every token a score. At each position, the model outputs one number per token in its vocabulary. nanochat’s vocabulary has 32,768 tokens, and Llama 3’s has 128,256. These numbers are called logits. They can be any size, positive or negative. Softmax turns them into probabilities in two steps. It raises e to each score, so every value is positive. Then it divides by the total, so the values add up to 1.

Line 4 turns one probability into a penalty. The data says which token actually came next. Call the probability the model gave that token p. The penalty is −log p. If p = 1, the penalty is 0. At p = 0.5 it is 0.69, and at p = 0.1 it is 2.30. As p gets close to 0, the penalty grows without limit, so the most expensive mistake is to be confident and wrong. F.cross_entropy computes the same number without forming p at all. It takes the log of the total and subtracts the right token’s score: log 189.72 − 5.0 = 0.2455. This form is also safer on a computer, because e raised to a large score can overflow.

The loss is the average penalty. Each training step averages the penalties over every token in its batch. Pre-training counts every token. The SFT step in post-training adds a mask so that only the answer counts. For the sentence above, the average is 1.3340. nanochat measures its validation loss the same way, over 80 × 524,288 tokens, which is about 42 million penalties averaged into one number.

A few tokens make up most of the loss. In that sentence, the three hardest tokens (“The”, “capital” and the “2” of 2.1 million) make up half of the total. The eight tokens the model gave 0.5 or more make up 15%. Nothing earlier in the text says the document will be about a capital, or that the population starts with 2, so no model can predict those tokens well. That is why the loss never reaches 0. Text has real uncertainty in it, and a model can only get close to that limit. The Chinchilla paper’s fit puts the limit at about 1.69 for its data.

loss    e^(-loss)
4.00    0.0183    about 1 in 55
3.28    0.0376    about 1 in 27    the modded-nanogpt speedrun target
2.00    0.1353    about 1 in 7
1.69    0.1845    about 1 in 5     the limit in the Chinchilla fit
1.00    0.3679    about 1 in 3

Bits per byte is the same number in other units. The loss uses the natural log, so its unit is nats per token. Divide by ln 2 = 0.6931 to get bits per token. Then divide by the average number of bytes in a token to get bits per byte. If a token covers about 4 bytes, the speedrun’s 3.28 becomes 3.28 ÷ 0.6931 ÷ 4 ≈ 1.18 bits per byte. The tokenizer step explains why runs report it: it does not depend on the tokenizer.

The pipeline is a sequence of filters, and together they throw away most of what comes in. A snapshot arrives as WARC files: raw HTTP responses, HTML and all. trafilatura pulls the main article out of the HTML. FineWeb's first finding was that extracting the text yourself from the WARC files works better than using the WET files, which hold text Common Crawl already extracted. The text comes out cleaner, and models trained on it do better.

Next comes language identification. A fastText classifier scores each page, and FineWeb keeps pages that score above 0.65 for English. Then come the quality rules inherited from Gopher and C4. They drop documents with too many repeated lines, too few words, no sentence-ending punctuation, too many symbols, the word "javascript", or "lorem ipsum". FineWeb added three rules of its own after testing each one in small training runs. They drop documents with too many short lines, too many lines without punctuation, or too much duplicated text.

Then deduplication. FineWeb uses MinHash over 5-grams, with 112 hash functions in 14 buckets of 8. That setup catches pairs of documents that are about 75 percent similar. Here FineWeb found something surprising. Deduplicating across all snapshots at once made models worse than deduplicating within each snapshot. Global dedup mostly removed good text that is copied all over the web, and kept junk that appears only once. Finally, personal data (PII): emails and IP addresses are replaced with placeholders.

Then the classifier

Hand-written rules get you to FineWeb. The next big jump was a learned quality classifier. For FineWeb-Edu, Llama-3-70B-Instruct scored about 460K pages on a 0 to 5 scale of "educational value". Those labels trained a small regression head on top of an embedding model. That small model then scored all 15T tokens, and only pages scoring 3 or above were kept. The 1.3T tokens that remained beat the full set on knowledge benchmarks. DCLM did the same job with a fastText classifier, trained to tell OpenHermes instruction data and good r/ELI5 answers apart from random web pages. DCLM kept the top 10 percent. Every serious dataset since uses some version of this: a big model labels a sample, then a cheap model scores everything.

datatrove is FineWeb's pipeline as code. A pipeline is a list of steps. Each step is a class that takes documents in and passes documents on. The real FineWeb file is longer, but this is its skeleton, with the real filter names and thresholds.

from datatrove.executor import SlurmPipelineExecutor
from datatrove.pipeline.readers import WarcReader
from datatrove.pipeline.extractors import Trafilatura
from datatrove.pipeline.filters import (URLFilter, LanguageFilter, GopherRepetitionFilter,
                                        GopherQualityFilter, C4QualityFilter, FineWebQualityFilter)
from datatrove.pipeline.writers import JsonlWriter

pipeline = [
    WarcReader("s3://commoncrawl/crawl-data/CC-MAIN-2024-10/segments/"),
    URLFilter(),                                       # blocklists + bad words in the url
    Trafilatura(favour_precision=True),                # html -> main text
    LanguageFilter(languages=["en"], language_threshold=0.65),
    GopherRepetitionFilter(),                          # dup lines / paragraphs / n-grams
    GopherQualityFilter(min_doc_words=50, max_doc_words=100_000),
    C4QualityFilter(filter_no_terminal_punct=False),   # fineweb kept most of c4, dropped this one
    FineWebQualityFilter(),                            # short-line and punctuation ratios
    JsonlWriter("s3://my-bucket/fineweb/filtered/"),
]
SlurmPipelineExecutor(pipeline, tasks=8000, time="10:00:00", partition="cpu").run()

Deduplication is a separate pipeline with four stages, because MinHash needs to see all the documents at once. First compute a signature for each document. Then sort the signatures into buckets. Then group near-duplicates into clusters. Finally, keep one document per cluster and remove the rest.

from datatrove.pipeline.dedup import (MinhashDedupSignature, MinhashDedupBuckets,
                                      MinhashDedupCluster, MinhashDedupFilter, MinhashConfig)

cfg = MinhashConfig(hash_config=HashConfig(hash_fc="sha1", precision=64),
                    num_buckets=14, hashes_per_bucket=8, n_grams=5)   # 112 hashes ≈ 75% jaccard
stage1 = [JsonlReader(filtered), MinhashDedupSignature(output_folder=sigs, config=cfg)]
stage2 = [MinhashDedupBuckets(input_folder=sigs, output_folder=buckets, config=cfg)]
stage3 = [MinhashDedupCluster(input_folder=buckets, output_folder=clusters, config=cfg)]
stage4 = [JsonlReader(filtered), MinhashDedupFilter(input_folder=clusters), JsonlWriter(deduped)]

And here is the classifier stage, cut down to its core. Score everything with a cheap model and keep the top scorers. This is FineWeb-Edu's version: a regression head on a small embedding model, trained once on the Llama-70B labels.

from transformers import AutoTokenizer, AutoModelForSequenceClassification
tok = AutoTokenizer.from_pretrained("HuggingFaceFW/fineweb-edu-classifier")
clf = AutoModelForSequenceClassification.from_pretrained("HuggingFaceFW/fineweb-edu-classifier")

def keep(doc):
    score = clf(**tok(doc["text"], truncation=True, max_length=512)).logits.item()
    doc["score"] = score                            # e.g. 3.375
    doc["int_score"] = round(max(0, min(score, 5)))  # -> 3
    return doc["int_score"] >= 3                    # fineweb-edu keeps 3, 4, 5 -> 1.3T tokens

{
  "id": "<urn:uuid:673b1bf6-2c30-40ae-992b-c387d00a836a>",
  "dump": "CC-MAIN-2013-20",
  "text": "No. 24; Updated March 2011 ... Parents are usually the first to recognize ...",
  "url": "https://www.aacap.org/AACAP/Families_and_Youth/Facts_for_Families/FFF-Guide/When-to-Seek-Help-for-Your-Child-024.aspx",
  "date": null,
  "file_path": "s3://commoncrawl/crawl-data/CC-MAIN-2013-20/segments/1368696381249/warc/CC-MAIN-20130516092621-00000-ip-10-60-113-184.ec2.internal.warc.gz",
  "language": "en", "language_score": 0.927742, "token_count": 755,
  "score": 3.375, "int_score": 3
}

What the other datasets add

// attributes/<tagger>/ file, same row order, holds the derived signals:// attributes/<tagger>/ file, same row order, holds the derived signals:// attributes/<tagger>/ file, same row order, holds the derived signals:als:
{"source": "...", "id": "...", "attributes": {"toxicity": 0.7,
  "olmo_mix_v1_taggers__ft_lang_id_en_paragraph_with_doc_score_v2__en": [[0, 300, 0.9], [300, 540, 0.3]]}}

What breaks at scale

• This is CPU work. The FineWeb recipe runs as thousands of Slurm CPU tasks reading from S3. The crawl adds up to petabytes across snapshots, so you stream it through and never store it all at once.
• Dedup is where memory runs out. MinHash signatures for 25B documents have to be shuffled across all machines so that similar ones end up together. That is why dedup takes four passes instead of one function call. DCLM's Bloom-filter approach (BFF, written in Rust, checking 13-gram overlap at the paragraph and document level) gives up exactness to do it in a single streaming pass.
• Decontamination. Benchmark questions leak onto the web. Every lab now removes text that shares 13-grams with its eval sets before training, and this still misses some.
• The classifier shapes what the model learns. FineWeb-Edu's labeling prompt asks for "educational" text, which tilts the corpus toward textbooks and away from fiction and dialogue. That choice shows up in the finished model, and it is why labs keep several classifiers and mix their outputs. Text extraction is the same kind of choice: trafilatura (FineWeb) and resiliparse (DCLM) produce different corpora from the same crawl.
• Licensing and opt-out. Some pipelines honor robots.txt and publisher opt-outs at crawl time, and others do not. Code datasets filter by license. Here the limits are set by lawyers rather than engineers.

Modern tokenizers are byte-level BPE. Start with the 256 possible bytes as your vocabulary. Count every adjacent pair in the training text. Merge the most frequent pair into a new token. Repeat until you hit the vocabulary size you chose. Because you started from bytes, any string in any language can be encoded; there is no "unknown" token.

Before BPE runs, the pre-tokenizer regex splits the text into chunks, and a merge can never cross from one chunk into another. A chunk is a word with its leading space, a run of up to three digits, a run of punctuation, or whitespace. That regex is why " generates" is one token and why "3,141" is three. GPT-4 introduced the three-digit rule so that all numbers are split the same way, and Llama 3 copied it almost verbatim.

Vocabulary size

The vocab size trades sequence length against the size of the embedding table. A bigger vocabulary means fewer tokens per document, which makes attention cheaper and fits more text into the context. But it also means a bigger embedding table and output head, and more rare tokens that each get less training.

The tokenizer also fixes the special tokens. A base model needs only one, the document separator. But the chat format taught in post-training needs more (markers for whose turn it is, brackets around tool calls), and the tokenizer cannot change later. So they are reserved up front. Llama 3 reserves 256 of them. nanochat defines nine, from <|bos|> to <|output_end|>, before pre-training starts.

Here are the regexes, verbatim. Each one describes what counts as a chunk, and a merge can never cross a chunk boundary.

# GPT-2 (tiktoken r50k)
'(?:[sdmt]|ll|ve|re)| ?\p{L}++| ?\p{N}++| ?[^\s\p{L}\p{N}]++|\s++$|\s+(?!\S)|\s

# GPT-4 (cl100k): digits capped at 3, case-insensitive contractions
'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}++|\p{N}{1,3}+| ?[^\s\p{L}\p{N}]++[\r\n]*+|\s++$|\s*[\r\n]|\s+(?!\S)|\s

# Llama 3 (llama-models/models/llama3/tokenizer.py, pat_str)
(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+

# nanochat (nanochat/tokenizer.py, SPLIT_PATTERN): digits capped at 2
'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,2}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+

Training a tokenizer is a short loop that any of us could write. This is the algorithm from Karpathy's minbpe. The Rust version in nanochat does the same thing with faster data structures.

def train_bpe(text: str, vocab_size: int):
    ids = list(text.encode("utf-8"))               # start from bytes: 256 base tokens
    merges = {}
    for new_id in range(256, vocab_size):
        counts = {}
        for a, b in zip(ids, ids[1:]):             # count adjacent pairs (inside regex chunks, in practice)
            counts[(a, b)] = counts.get((a, b), 0) + 1
        pair = max(counts, key=counts.get)         # the most frequent pair becomes one token
        merges[pair] = new_id
        ids = merge(ids, pair, new_id)             # replace every occurrence
    return merges                                  # the tokenizer *is* this ordered list of merges

nanochat's script is the same idea at full size. It caps the total training text and the length of each document, trains in Rust, then hands the merges to tiktoken for fast encoding.

# scripts/tok_train.py: defaults --max-chars 2_000_000_000 --doc-cap 10_000 --vocab-size 32768
tokenizer = RustBPETokenizer.train_from_iterator(text_iter, args.vocab_size)   # rustbpe, GPT-4-style split

# nanochat/tokenizer.py: wrap the learned merges in tiktoken for encode_ordinary_batch speed
enc = tiktoken.Encoding(name="rustbpe", pat_str=SPLIT_PATTERN,
                        mergeable_ranks=mergeable_ranks, special_tokens=SPECIAL_TOKENS)
SPECIAL_TOKENS = ["<|bos|>", "<|user_start|>", "<|user_end|>", "<|assistant_start|>", "<|assistant_end|>",
                  "<|python_start|>", "<|python_end|>", "<|output_start|>", "<|output_end|>"]

What production adds

• Pad the vocabulary. The embedding table and output head are padded to a multiple of 64 or 128 rows, because the GPU's tensor cores multiply matrices fastest when the sizes divide evenly into their tiles. nanochat uses pad_vocab_size_to=64; OLMo pads to 128. The model never outputs the extra rows.
• Measure in bits per byte, not loss per token. Cross-entropy per token (the loss, explained under the map) depends on the tokenizer, since bigger tokens mean fewer but harder predictions. Divide the total loss by the number of bytes the tokens cover, and runs with different vocabularies become comparable. nanochat reports val_bpb for this reason.
• How to split digits. GPT-4 and Llama 3 split numbers into groups of three digits, nanochat into groups of two, Qwen into single digits. Each is a guess about what will help the model do arithmetic. Nobody agrees yet, and there are only experiments to go on.
• Byte-level vs byte fallback. Llama 2 used SentencePiece with byte_fallback: its tokens are built from characters, and it falls back to raw bytes only for characters it does not know. Llama 3 switched to tiktoken's byte-level BPE, and so did nearly everyone else.
• Tokens per word in other languages. A vocabulary trained mostly on English spends three to five tokens per word on Thai or Hindi. Qwen3's 151K and Gemma's 262K vocabularies are sized to bring that number down. Apertus, trained on 1,800 languages, makes covering all of them well a main design goal of its tokenizer.
• The tokenizer depends on its training text. nanochat trains it on the first 2B characters of the same corpus it will pre-train on. Train it on the wrong mix of text, and code or math will split into poor tokens for the life of the model.

The 2023 default, still used in nanoGPT and in torchtitan's simplest loader, is "concatenate with EOS, cut every T tokens". It wastes no space, but it cuts documents mid-sentence and lets attention reach from one document into another. Two improvements are now standard.

Best-fit packing treats each row as a bin to fill. nanochat's loader starts every row with BOS, then keeps placing the largest document that still fits. When nothing fits, it crops one document to fill the row exactly. The result: no padding, every token can see the start of its own document, and about 35 percent of tokens get cropped off at T = 2048. The paper behind it is "Fewer Truncations Improve Language Modeling" (Ding et al., 2024).

Document masking stops attention from crossing between documents. The attention mask becomes block-diagonal within the row, so a token never looks back past the start of its own document. This needs an attention kernel that handles variable-length sequences: FlashAttention's varlen interface with cu_seqlens, or FlexAttention with a document mask. OLMo-core passes doc_lens (the document lengths) through the whole stack for this.

Mixing, and the three stages

The corpus is made of dozens of sources, each with a weight. The mixture says how many tokens come from each source. A source that is scarce and valuable can be repeated for more than one epoch (OLMo-core's max_repetition_ratio sets the limit). SmolLM3's config lists about 45 sources: FineWeb-Edu, DCLM, papers, Wikipedia, StackExchange, 14 languages of FineWeb2, two math sets, 15 languages of The Stack, pull requests, Kaggle notebooks, GitHub issues.

The weights also change over the run. In 2025 and 2026 the norm is three stages. First comes a long general stage at short context. Second, a shorter "midtraining" stage gives more weight to math, code and higher-quality text while the learning rate decays to zero. Third, a brief long-context stage trains at 32K to 64K tokens.

nanochat's loader tokenizes parquet shards on the fly and packs with best-fit. This is the core of tokenizing_distributed_data_loader_with_state_bos_bestfit, trimmed.

row_capacity = T + 1                                  # +1 so inputs and targets can both be T long
token_lists = tokenizer.encode(doc_batch, prepend=bos_token, num_threads=tokenizer_threads)
doc_buffer.extend(token_lists)
for row_idx in range(B):
    pos = 0
    while pos < row_capacity:
        remaining = row_capacity - pos
        best_idx, best_len = -1, 0
        for i, doc in enumerate(doc_buffer):          # largest document that fits entirely
            if len(doc) <= remaining and len(doc) > best_len:
                best_idx, best_len = i, len(doc)
        if best_idx >= 0:
            doc = doc_buffer.pop(best_idx)
            row_buffer[row_idx, pos:pos + len(doc)] = torch.tensor(doc); pos += len(doc)
        else:                                         # nothing fits: crop the shortest to fill exactly
            doc = doc_buffer.pop(shortest_idx)
            row_buffer[row_idx, pos:pos + remaining] = torch.tensor(doc[:remaining]); pos += remaining
inputs, targets = row_buffer[:, :-1], row_buffer[:, 1:]   # (B, T) each; targets are inputs shifted left

Here is document masking done two ways. modded-nanogpt finds the BOS positions in the stream and passes those document boundaries to FlashAttention-3. torchtitan builds the same mask as a FlexAttention block mask.

# modded-nanogpt: boundaries as cumulative sequence lengths, batch size 1
# seqlens = [0, 5, 14, 16, ...]  cumulative document starts within the packed stream
y = flash_attn_varlen_func(q[0], k[0], v[0], cu_seqlens_q=seqlens, cu_seqlens_k=seqlens,
                           max_seqlen_q=max_len, max_seqlen_k=max_len,
                           causal=True, window_size=(bm_size, 0))          # sliding window too

# torchtitan: the same mask as a predicate for FlexAttention
def document_causal(b, h, q_idx, kv_idx):
    return (q_idx >= kv_idx) & (doc_id[b, q_idx] == doc_id[b, kv_idx])
block_mask = create_block_mask(document_causal, B, None, T, T)
y = flex_attention(q, k, v, block_mask=block_mask)

The mix is set in configuration files rather than written as code. OLMo-core's source mixture is a list of sources with target ratios. The official mix manifest is 951 lines long, with one file path per line.

SourceMixtureConfig(source_name="dclm",      target_ratio=0.55, paths=[...], max_repetition_ratio=1.0)
SourceMixtureConfig(source_name="stack_edu", target_ratio=0.15, paths=[...], max_repetition_ratio=2.0)  # up-sample: two epochs
SourceMixtureConfig(source_name="math",      target_ratio=0.05, paths=[...])

# data/mixes/OLMo-mix-0925-official.txt · one line per tokenized shard
common_crawl_adult_content,preprocessed/dolma3-0625/v0.1-official/{TOKENIZER}/common_crawl/adult_content/000000.npy

What the tokens look like on disk

Where it gets hard

• Cropping vs padding. nanochat accepts losing 35 percent of tokens to cropping because web data is plentiful. With scarce data (a specialist domain, a small language) you keep the older loader and accept that some documents get cut in the middle.
• A simple mask is slow. A dense (T×T) mask passed to standard attention costs as much as full attention, plus the cost of the mask. The varlen and flex kernels skip the masked-out blocks entirely, which is why they exist.
• Long context is mostly a RoPE change. Raise the rotary base (10K to 500K for Llama 3 and OLMo, 1M for Qwen3, 1.5M then 5M for SmolLM3), train briefly at the longer length, and stretch further at inference with YaRN. The cost at training time is context parallelism: OLMo 3 splits 65K-token sequences across 8 GPUs.
• The published base model is the checkpoint after the anneal. Midtraining is where the learning rate reaches zero and math and code get more weight. So the "base model" you download is not the model from the long constant-learning-rate stage. Labs also use annealing runs to compare data: branch from a checkpoint in the stable stage, anneal on each candidate mix, and compare the results.
• Epochs vs rephrasing. Kimi K2 reuses its most valuable text by having an LLM rephrase it instead of repeating it word for word. The evidence is that repeats stop helping after a couple of epochs.
• Resume must be deterministic. The loader's position is saved in the checkpoint along with the model, so a restart continues from the same document. nanochat's resume is approximate (it skips ahead one row group to avoid repeating data). OLMo's is exact, because its index records where every document is.

Read the middle column from top to bottom. RMSNorm replaced LayerNorm. It only divides each vector by its root-mean-square, with no mean subtraction and no bias, and nanochat drops the learnable gain too. It is cheaper and works as well.

RoPE replaced learned position embeddings. Instead of adding a position vector to the input, it rotates the query and key vectors by an angle that grows with position. That way the dot product between two tokens depends only on how far apart they are. RoPE can generalize past the training length, which learned positions never could. It is also why extending the context comes down to changing one number, the base frequency.

GQA (grouped-query attention) shrinks the KV cache, the keys and values kept for past tokens during generation. In Llama 3 8B, 32 query heads share 8 key-value heads; in Qwen3-32B, 64 share 8. The cache memory at inference drops fourfold or more, for a tiny cost in quality. QK-norm applies RMSNorm to queries and keys before their dot product. That keeps attention logits from growing too large, and it is the most common fix for training instability in the 2025 recipes (OLMo 2 and 3, Qwen3, nanochat).

SwiGLU replaced the GELU MLP. It uses three matrices instead of two: a gate that is multiplied into an up-projection, then a down-projection. The hidden width drops from 4× to 8/3× the model width, so the parameter count stays the same. The speedrun projects use ReLU² instead, which needs one matrix fewer and works as well at their scale. Biases are gone everywhere, and so is dropout. The embedding table and the output head are usually separate matrices (untied).

Two newer additions. First, sliding-window layers. In each group of four layers, three attend only to the last 768 or 4,096 tokens and the fourth attends to everything (nanochat's pattern is "SSSL", OLMo 3's is [4096, 4096, 4096, -1]). Most of the cost of attention goes away, and the full layers still see the whole context. Second, a logit softcap: 15 · tanh(logits / 15) on the output, from Gemma 2. It keeps every final logit between −15 and 15, so none can grow without limit.

nanochat's gpt.py shows the 2026 dense block in 555 lines. Here is the attention forward pass, lightly trimmed. Notice what is missing: no positional embedding table, no bias, no dropout.

def forward(self, x, ve, cos_sin, window_size, kv_cache):
    B, T, C = x.size()
    q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
    k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)      # n_kv_head < n_head: GQA
    v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
    if ve is not None:                                              # value embeddings on alternate layers
        gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))
        v = v + gate.unsqueeze(-1) * ve
    cos, sin = cos_sin
    q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin)   # RoPE
    q, k = norm(q), norm(k)                                         # QK-norm
    y = flash_attn.flash_attn_func(q, k, v, causal=True, window_size=window_size)  # (768, 0) or (T, 0)
    return self.c_proj(y.contiguous().view(B, T, -1))

class MLP(nn.Module):
    def forward(self, x):
        return self.c_proj(F.relu(self.c_fc(x)).square())           # relu²: one matrix fewer than SwiGLU

class Block(nn.Module):
    def forward(self, x, ve, cos_sin, window_size, kv_cache):
        x = x + self.attn(norm(x), ve, cos_sin, window_size, kv_cache)   # pre-norm, residual
        x = x + self.mlp(norm(x))
        return x

def norm(x):
    return F.rms_norm(x, (x.size(-1),))                             # no learnable gain

Here is the SwiGLU version from the Llama reference code, so you can see where the 8/3 comes from.

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, multiple_of, ffn_dim_multiplier):
        hidden_dim = int(2 * hidden_dim / 3)                          # 4d -> 8d/3: same params as a 2-matrix mlp
        if ffn_dim_multiplier is not None:
            hidden_dim = int(ffn_dim_multiplier * hidden_dim)         # llama 3.1 8b: 1.3 -> 14,336
        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
        self.w1 = Linear(dim, hidden_dim, bias=False)                 # gate
        self.w2 = Linear(hidden_dim, dim, bias=False)                 # down
        self.w3 = Linear(dim, hidden_dim, bias=False)                 # up
    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

What the frontier models add

MLA, multi-head latent attention, is DeepSeek's way to shrink the KV cache. It compresses each token's keys and values into one 512-dimensional vector, stores only that, and expands it again when needed. A separate 64-dimensional "rope key" carries the position, because the RoPE rotation cannot pass through the compression. Queries get their own low-rank path (rank 1,536). In absorb mode, the expansion matrix is folded into the attention-score computation, so inference never builds the full keys at all. DeepSeek-V3 and Kimi K2 both use MLA. Kimi K3 adds a gate to it and interleaves it with linear-attention layers (KDA): 69 linear layers to 24 gated-MLA layers.

MTP, multi-token prediction, adds one small extra block that predicts the token after next from the main model's hidden state. It has its own cross-entropy loss, weighted 0.3 for the first 10T tokens and 0.1 after that. It gives the model more to learn from each position, and at inference it can serve as the draft model for speculative decoding. torchtitan's models/deepseek_v3/MTP.md documents the block: enorm, hnorm, eh_proj, one transformer block, and a shared output head.

DeepSeek V4 (April 2026) cuts attention memory further: sliding windows plus heavily compressed attention and compressed sparse attention, with a 1M context. Kimi K3 is a hybrid with linear attention. Nemotron 3 is a hybrid Mamba-Transformer MoE. They all point the same way: in the newest models, full attention over the whole sequence is the exception.

What the speedrun adds, and a warning

modded-nanogpt's file has value embeddings, U-net style skip connections between early and late layers (now "MUDD" skips), a bigram hash embedding, paired-head attention, learnable cross-sequence attention, and an FP8 output head. Its README warns that "some methods used in the speedrun are unlikely to scale, particularly those imposing additional network structure, such as logit softcapping."

Stability choices, side by side

The router is a single linear layer. It maps the token's hidden state to one score per expert (256 of them in V3), squashes each score into the range 0 to 1 with a sigmoid, and picks the top-k (8). Those experts process the token, and their outputs are added up, each weighted by its score (rescaled so the weights sum to 1). One shared expert processes every token, so common knowledge does not have to be copied into each routed expert. Kimi K2 keeps the same design with 384 experts. Kimi K3 has 896, picks the top 16, and has two shared experts. Qwen3-MoE has 128, top-8, and no shared expert.

How the router learns

MoE does not change what the model is trained to do. The input is still tokens, the output is still a probability for every possible next token, and the loss is still −log p of the token that actually came next. MoE only changes what happens inside each block. The router is trained by that same loss, from scratch, like every other weight in the model. Nobody tells it which expert should handle what.

The router learns through the mixing weights. If the prediction would have been better with a bit more of expert 2’s output, the gradient raises expert 2’s score for tokens that look like this one. So nothing assigns topics to experts. They specialize in whatever lowers the loss. When Mistral looked at how Mixtral routes tokens, they found no clear topic experts, only patterns at the level of tokens and syntax.

The router is trained here, once, and post-training inherits it. It causes a problem there too: the engine that generates answers and the trainer can send the same token to different experts, which the GSPO fix in post-training works around.

The main way this goes wrong is load imbalance. Nothing stops the router from sending every token to the same three experts. Those three then act as a small dense model while the other 253 sit idle. The classic fix is an extra loss term that penalizes uneven use of experts, but that term pulls against the main language-modeling loss. DeepSeek-V3 avoids the extra loss. Instead it adds a per-expert bias to the scores, used only to choose the experts. After each step, an overloaded expert's bias goes down by a small constant and an underloaded expert's bias goes up. The weights used to mix the expert outputs never include the bias. Toggle it in the diagram above.

The memory-versus-compute math

Compare the total and active numbers. The compute per token grows with active parameters, so in compute V3 trains and serves like a 37B dense model. Memory grows with total parameters, so the cluster has to hold the weights of a 671B model. MoE is a bet on that trade: parameters are a cheap place to store knowledge, and compute is the scarce resource. Kimi K2's report fits a "sparsity scaling law" to pick 384/8 = 48 as the ratio of total experts to active experts.

DeepSeek-V3's inference code is the clearest reference I have found. Here is the gate, verbatim except for the comments. It has every piece: sigmoid scores, the balancing bias, node-limited routing, top-k, and renormalization.

def forward(self, x):                                    # class Gate, inference/model.py
    scores = linear(x, self.weight)                                  # (tokens, n_experts)
    scores = scores.softmax(-1, dtype=torch.float32) if self.score_func == "softmax" else scores.sigmoid()
    original_scores = scores
    if self.bias is not None:
        scores = scores + self.bias                                  # balancing bias: affects *which*, not *how much*
    if self.n_groups > 1:                                            # group-limited routing: pick 4 of 8 groups first
        scores = scores.view(x.size(0), self.n_groups, -1)
        group_scores = scores.amax(-1) if self.bias is None else scores.topk(2, -1)[0].sum(-1)
        indices = group_scores.topk(self.topk_groups, -1)[1]
        mask = scores.new_ones(x.size(0), self.n_groups, dtype=bool).scatter_(1, indices, False)
        scores = scores.masked_fill_(mask.unsqueeze(-1), float("-inf")).flatten(1)
    indices = torch.topk(scores, self.topk, dim=-1)[1]               # the 8 experts for each token
    weights = original_scores.gather(1, indices)                     # mixing weights come from the *unbiased* scores
    if self.score_func == "sigmoid":
        weights /= weights.sum(dim=-1, keepdim=True)                 # renormalize to 1
    weights *= self.route_scale                                      # 2.5 in v3
    return weights.type_as(x), indices

The MoE layer itself is a loop over experts. It is simple and correct, and it is the version to understand before the fast one.

def forward(self, x):                                    # class MoE
    shape = x.size()
    x = x.view(-1, self.dim)
    weights, indices = self.gate(x)                      # (N, k), (N, k)
    y = torch.zeros_like(x)
    counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist()
    for i in range(self.experts_start_idx, self.experts_end_idx):    # only this rank's experts
        if counts[i] == 0:
            continue
        expert = self.experts[i]
        idx, top = torch.where(indices == i)              # which tokens chose expert i, and in which slot
        y[idx] += expert(x[idx]) * weights[idx, top, None]
    z = self.shared_experts(x)                            # every token, always
    if world_size > 1:
        dist.all_reduce(y)                                # sum the partial results from every rank's experts
    return (y + z).view(shape)

The bias update lives in the training loop, not in the model, and it comes down to one line. γ was 0.001 for the first 14.3T tokens of V3 and 0 for the rest.

load = torch.bincount(indices.flatten(), minlength=n_experts).float()   # tokens per expert this step
gate.bias += gamma * torch.sign(load.mean() - load)                        # overloaded: down, underloaded: up

Production code replaces the Python loop with a sort and a grouped GEMM, one matrix multiply over many groups of different sizes. Sort the tokens by expert, run one batched matmul over all the groups, then put the tokens back in their original order. torchtitan's GroupedExperts does this with torch._grouped_mm; MegaBlocks did it first with block-sparse kernels.

order = indices.flatten().argsort()                       # tokens grouped by expert id
x_sorted = x.repeat_interleave(k, 0)[order]               # each token appears k times, once per chosen expert
offsets = counts.cumsum(0)                                # where each expert's group ends
h = torch._grouped_mm(x_sorted, w1, offs=offsets)         # one kernel, all experts, no padding
h = F.silu(torch._grouped_mm(x_sorted, w_gate, offs=offsets)) * h
out_sorted = torch._grouped_mm(h, w2, offs=offsets)
y = torch.zeros_like(x).index_add_(0, token_of[order], out_sorted * weights.flatten()[order, None])

Expert parallelism

256 experts do not fit on one GPU, so they are spread across an expert-parallel group of GPUs, and each token is sent to the GPUs that hold its experts. That takes one all-to-all exchange before the experts (dispatch) and another after (combine). DeepSeek-V3 trained with 64-way EP across 8 nodes. Its node-limited routing exists so that a token's 8 experts sit on at most 4 nodes, which limits the traffic between nodes. DeepEP is DeepSeek's open kernel for the all-to-all. Kimi K2 used 16-way EP with 16-way pipeline parallelism.

EP and FSDP do not fit together neatly. The experts are split one way (by expert) and everything else another way (by parameter). torchtitan keeps a separate "edp" mesh for the expert weights and prefetches expert all-gathers around the all-to-all.

Where MoE training breaks

• Router collapse. Without balancing, a few experts get chosen early and keep getting chosen, and the rest never learn. The Switch Transformer loss is α · N · Σᵢ fᵢ · Pᵢ (the fraction of tokens routed to expert i, times the mean router probability for i). DeepSeek-V3 keeps a tiny per-sequence version at α = 0.0001 alongside the bias. Qwen3 measures balance over the global batch instead of each micro-batch, so an expert that specializes in one domain is not penalized.
• Router in fp32. The scores are computed in float32 even in a bf16 or FP8 run. Tiny differences in score decide which experts are chosen, and low precision makes those differences noisy. DeepSeek keeps the gating modules out of FP8 entirely.
• Capacity factor vs dropless. Older MoEs capped how many tokens each expert could take and dropped the overflow. Dropless MoE (MegaBlocks) processes every token, using groups of different sizes, which grouped GEMM makes possible.
• Upcycling. Copy a dense model's MLP into every expert and continue training. Qwen and MiniCPM did this. It saves starting from scratch, but the experts begin identical and take a while to become different from each other.
• Fine-tuning fragility. A small post-training dataset barely reaches most experts. Kimi K2's advice is to fine-tune a MoE that was pretrained with Muon using Muon again. How well LoRA works on MoE layers is still unsettled in published work.
• Serving is a different problem. A model with 37B active parameters still needs all 671B held in memory. Batching many requests together is what makes it economical to serve, and a single user's latency gets no benefit from it.

Initialization

Nearly every lab initializes every matrix from a truncated normal with standard deviation 0.02 (OLMo's InitMethod.normal, torchtitan's default), and scales the output projections down with depth so that the residual stream does not grow layer by layer. The speedrun projects go further and start the output projections of attention and the MLP at zero, so every block begins by passing its input through unchanged. nanochat's table is worth reading in full, because every line is a deliberate choice.

wte (embedding):     normal, std=0.8   (then RMS-normed, so the scale is free)
lm_head:             normal, std=0.001
attn.c_q, c_k, c_v:  uniform, std=1/sqrt(n_embd)   (uniform avoids outliers)
attn.c_proj:         zeros
mlp.c_fc:            uniform, std=0.4/sqrt(n_embd)
mlp.c_proj:          zeros
resid_lambdas:       1.15 -> 1.05 across depth     (per-layer residual scale)
x0_lambdas:          0.20 -> 0.05 across depth     (blend the input embedding back in)

AdamW, the boring consensus

β₁ = 0.9, β₂ = 0.95, weight decay 0.1, gradient clipping at 1.0, and no weight decay on the embeddings. That line describes Llama 3, DeepSeek-V3, OLMo 3, SmolLM3 and Qwen3. Peak learning rates: 3e-4 for OLMo 3 7B, 2e-4 for SmolLM3 and Kimi K2, 2.2e-4 for DeepSeek-V3. Warmup is about 2,000 steps. Labs differ in the schedule and the batch size, and hardly at all in the optimizer.

Muon

Muon is the exception. It is used for the 2-D weight matrices inside the blocks. It takes the gradient, averaged with momentum, and orthogonalizes it: it replaces the update matrix with the nearest matrix whose singular values are all one. Five Newton-Schulz iterations in bf16 compute this. As a result, every direction in the update gets the same step size, instead of a few dominant directions taking almost all of it. Embeddings, the output head and scalar parameters stay on AdamW. Muon won the speedrun in 2024. Moonshot showed it scales ("Muon is Scalable": 16B parameters on 5.7T tokens, with about half the FLOPs of AdamW for the same loss), and Kimi K2 trained a trillion-parameter model with it. MuonClip is K2's addition. After each step, if any attention head's largest logit is above a threshold τ, that head's query and key weights are scaled down. K2 reports 15.5T tokens without a single loss spike.

The schedule, the batch, the budget

Two kinds of schedule are in use. One is cosine decay to about 10 percent of the peak (Llama 2 and 3, OLMo 3, Kimi K3). The other is WSD: warmup, then stable, then decay (nanochat, an option in OLMo-core, SmolLM3, and DeepSeek-V3 and Kimi K2 in a multi-phase form). WSD's flat middle is what makes annealing experiments and run extensions possible. A batch-size ramp is standard too. DeepSeek-V3 grows its batch from 3,072 to 15,360 sequences over the first 469B tokens. nanochat settled on 1M tokens per step.

Then there is the budget. Chinchilla's rule of 20 tokens per parameter is now treated as a minimum. Labs train far past it, because inference cost depends on the number of parameters and not on how many tokens went into training. A smaller model trained longer is cheaper to serve.

Here is Muon, verbatim from the modded-nanogpt README. The three constants are the coefficients of a quintic polynomial, tuned so the iteration converges quickly in bf16.

def zeroth_power_via_newtonschulz5(G, steps=5, eps=1e-7):
    a, b, c = (3.4445, -4.7750, 2.0315)
    X = G.bfloat16() / (G.norm() + eps)                  # normalize so the iteration converges
    if G.size(0) > G.size(1): X = X.T
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * A @ A
        X = a * X + B @ X                                # X -> polynomial in X X^T, pushes singular values to 1
    if G.size(0) > G.size(1): X = X.T
    return X.to(G.dtype)

# the step, per matrix parameter
buf.mul_(momentum).add_(grad)                           # nesterov momentum
g = grad.add(buf, alpha=momentum)
g = zeroth_power_via_newtonschulz5(g)                    # orthogonalize
p.add_(g, alpha=-lr * max(1, p.size(0) / p.size(1)) ** 0.5)   # scale by aspect ratio

nanochat's optimizer setup shows the split: which parameters get Muon, which get AdamW, and how the AdamW learning rates scale with width so that one config works for every depth.

dmodel_lr_scale = (model_dim / 768) ** -0.5              # AdamW lr ∝ 1/√d_model, tuned at 768
param_groups = [
    dict(kind='adamw', params=lm_head_params,   lr=0.004 * dmodel_lr_scale, betas=(0.8, 0.96),  weight_decay=0.01),
    dict(kind='adamw', params=embedding_params, lr=0.2   * dmodel_lr_scale, betas=(0.8, 0.995), weight_decay=0.001),
    dict(kind='adamw', params=[resid_lambdas],  lr=0.005, betas=(0.8, 0.95), weight_decay=0.05),
    dict(kind='adamw', params=[x0_lambdas],     lr=0.5,   betas=(0.96, 0.95), weight_decay=0.0),
]
for shape in sorted({p.shape for p in matrix_params}):   # every 2-D weight inside the blocks
    param_groups.append(dict(kind='muon', params=[p for p in matrix_params if p.shape == shape],
                             lr=0.02, momentum=0.95, ns_steps=5, beta2=0.9, weight_decay=weight_decay))
optimizer = MuonAdamW(param_groups)

The schedule takes about twenty lines. This is nanochat's: a WSD with a long linear warmdown. The Muon momentum has its own schedule alongside it.

def get_lr_multiplier(it):                               # scripts/base_train.py
    warmup_iters = args.warmup_steps                     # 40
    warmdown_iters = round(args.warmdown_ratio * num_iterations)   # 0.65 of the run
    if it < warmup_iters:
        return (it + 1) / warmup_iters
    elif it <= num_iterations - warmdown_iters:
        return 1.0
    else:
        progress = (num_iterations - it) / warmdown_iters
        return progress * 1.0 + (1 - progress) * args.final_lr_frac   # down to 0.05

def get_muon_momentum(it):                               # 0.85 -> 0.97 over 400 steps, back to 0.90 in the warmdown
    ...

num_iterations = (target_param_data_ratio * num_scaling_params) // total_batch_size   # the budget, from one dial

Here is the core of MuonClip, from the Kimi K2 report. It runs after the optimizer step, so the forward and backward passes of the current step are unchanged.

for h in heads:
    S_max = (1 / sqrt(d)) * max_over_batch(Q_h @ K_h.T)          # this head's largest attention logit
    if S_max > tau:                                              # tau = 100
        gamma = tau / S_max
        W_q[h] *= gamma ** alpha                                 # alpha = 0.5: split the shrink between q and k
        W_k[h] *= gamma ** (1 - alpha)

What the experiment logs say

• Weight decay should shrink as the model gets wider. nanochat's log found the best value goes as 1/width² (0.22 at depth 12, 0.08 at depth 20). nanochat also lowers it linearly to zero over the run.
• Gradient clipping can be dropped. With Muon and QK-norm, nanochat removed clipping ("grad norm never exceeds 1.0 naturally"), and modded-nanogpt does not clip either. Every production run with AdamW still clips at 1.0.
• Muon keeps changing. nanochat replaced Newton-Schulz with Polar Express coefficients, added row equilibration, rescales each update to a target Frobenius norm, and uses NorMuon's factored second moments and "cautious" weight decay (decay only where the update and the weight have the same sign). Kimi K3 orthogonalizes per attention head. Megatron-Core added Muon support in 2026.
• z-loss vs softcap. OLMo 3 adds 1e-5 · log²(Σ exp logits) to keep the output logits from drifting; DCLM used 5e-6. The speedrun projects use a tanh softcap instead. Both fix the same instability.
• Nobody uses muP in production. muP is a way of scaling a model so that settings tuned on a small version carry over to a big one. Every recipe I verified (Llama 3, DeepSeek-V3, OLMo 3, Kimi K2 and K3, Qwen3, SmolLM3) instead picks its learning rate and batch size by fitting scaling laws to runs of small models. Qwen3's report says it "sets the predicted optimal learning rate and batch size for each dense or MoE model". nanochat derives everything from --depth.
• The full DeepSeek-V3 schedule, for the record: 2K warmup steps to 2.2e-4, constant until 10T tokens, cosine to 2.2e-5 over the next 4.3T, constant at 2.2e-5 for 333B, then 7.3e-6 for the final 167B.

The loop, with accumulation

On one node, the loop looks like nanochat's. It runs a few forward-backward passes and adds up their gradients (gradient accumulation, so a 1M-token batch fits in memory), then takes one optimizer step. Activations are in bf16, the master weights and optimizer state are in fp32, and torch.compile fuses operations together. The multi-node version has the same shape, with four kinds of splitting added on top.

Data parallel, then FSDP

Plain data parallelism copies the whole model to every GPU, gives each GPU different data, and averages the gradients. It stops working when the model does not fit on one GPU. FSDP (fully sharded data parallel, called ZeRO-3 in DeepSpeed) keeps only a slice of every parameter on each GPU. Right before a layer runs, the GPUs send each other their slices so each has the full weights (an all-gather). Right after, they sum the gradients and each keeps only its own slice (a reduce-scatter). Memory per GPU drops by a factor equal to the number of GPUs. The cost is extra communication, which has to happen while the GPUs compute so it does not slow them down. FSDP2 in PyTorch is one function, fully_shard, called once per block and once on the root model. OLMo 3 trained its 7B on 512 H100s with only this, in a hybrid form (full copies across nodes, sharded within each node).

Tensor parallel

Split each weight matrix across the GPUs of one node. The first matmul of a block is split by columns and the second by rows, so the GPUs only need to combine their results once per sub-block (one all-reduce). Sequence parallel extends the split to the norms and dropout between blocks. Tensor parallel stays inside one node because every layer needs the fast NVLink connection between the GPUs. Llama 3 405B uses 8-way TP.

Pipeline parallel

Put layers 0 to 15 on one set of GPUs and 16 to 31 on another. The batch is cut into micro-batches that flow through the stages, and the schedule (1F1B, interleaved, zero-bubble, DeepSeek's DualPipe) decides how much of the time a stage sits idle waiting for its neighbor. Kimi K2 used 16-way pipeline with virtual stages; DeepSeek-V3, 16-way with DualPipe.

Context and expert parallel

Context parallel splits the sequence itself across GPUs and passes keys and values around a ring. That is how OLMo 3 trains on 65K-token sequences, split 8 ways. Expert parallel puts different experts on different GPUs and moves tokens to them with an all-to-all; DeepSeek-V3 ran 64-way EP across 8 nodes. torchtitan's rule for combining them: the dense part of the model is sharded over (dp × cp × tp) and the sparse part over (efsdp × ep), and both cover the same set of GPUs.

Precision

The default is bf16 for activations and the matmuls, and fp32 for the master weights, the optimizer moments and the gradient reduction. FP8 goes one step further, but only on the big linear layers. DeepSeek-V3's recipe quantizes to E4M3 with one scale for each 1×128 tile of activations and each 128×128 block of weights. Every 128 elements, it adds the partial sums into an fp32 accumulator. Embeddings, the output head, the MoE router, norms and attention stay in bf16 or fp32. Llama 4 trained in FP8; Nemotron 3 went to NVFP4; Kimi K3 uses MXFP4 weights with quantization-aware training from SFT onward, so the served model matches the trained one.

Here is nanochat's loop, verbatim but trimmed. The next batch is fetched while the GPU is busy with the backward pass.

for micro_step in range(grad_accum_steps):
    loss = model(x, y)
    train_loss = loss.detach()
    loss = loss / grad_accum_steps
    loss.backward()
    x, y, dataloader_state_dict = next(train_loader)      # prefetch next batch while GPU is busy
lrm = get_lr_multiplier(step)
muon_momentum = get_muon_momentum(step)
for group in optimizer.param_groups:
    group["lr"] = group["initial_lr"] * lrm
    if group["kind"] == "muon":
        group["momentum"] = muon_momentum
optimizer.step()                                           # gradient sync happens inside, no DDP
model.zero_grad(set_to_none=True)

# grad_accum_steps = total_batch_size // (device_batch_size * max_seq_len * world_size)

This is all FSDP2 takes. Each call wraps one unit. Before the forward and backward passes, the unit all-gathers its bf16 parameters. Afterward, it reduce-scatters the bf16 gradients into fp32 sharded gradients.

from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy

mp_policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
fsdp_config = {"mesh": dp_mesh, "mp_policy": mp_policy}
for layer_id, transformer_block in model.layers.items():
    fully_shard(transformer_block, **fsdp_config, reshard_after_forward=reshard_after_forward)
fully_shard(model, **fsdp_config)                          # the root: embeddings, norm, head

torchtitan's step adds two things a cluster needs. It computes a gradient norm that is correct across pipeline stages and expert groups. And if the loss or gradient norm is NaN or infinite, it stops training before the optimizer can apply the bad update.

grad_norm = dist_utils.clip_grad_norm_(
    [p for m in self.model_parts for p in m.parameters()],
    self.config.training.max_norm, foreach=True,
    pp_mesh=parallel_dims.get_optional_mesh("pp"), ep_enabled=parallel_dims.ep_enabled)
step_is_finite = loss_is_finite.logical_and(torch.isfinite(grad_norm).all())
torch._assert_async(step_is_finite,
    "Loss or gradient norm is not finite on at least one rank at "
    f"step {self.step}. Stopping training before the optimizer update.")
self.checkpointer.maybe_wait_for_staging()               # async checkpoint copy must finish before params change
self.optimizers.step()
self.lr_schedulers.step()

The parallel layout is set in configuration. torchtitan dropped its TOML files in 2026, and configs are Python now. This is the registered DeepSeek-V3 671B entry, trimmed.

# torchtitan/models/deepseek_v3/config_registry.py · deepseek_v3_671b
optimizer=default_adamw(lr=2.2e-4),
lr_scheduler=LRSchedulersContainer.Config(warmup_steps=2000, decay_ratio=0.8, decay_type="cosine", min_lr_factor=0.1),
training=TrainingConfig(num_tokens_per_microbatch_per_dp_rank=4 * model_spec.max_context_length, steps=10000),
parallelism=ParallelismConfig(pipeline_parallel_schedule="Interleaved1F1B", expert_parallel_degree=2),
checkpoint=CheckpointManager.Config(interval=500),
activation_checkpoint=SelectiveAC.Config(),
compile=CompileConfig(enable=True, components=["loss"]),

# ParallelismConfig fields (defaults): data_parallel_replicate_degree=1, data_parallel_shard_degree=-1,
#   tensor_parallel_degree=1, pipeline_parallel_degree=1, pipeline_parallel_schedule="1F1B",
#   context_parallel_degree=1, expert_parallel_degree=1
# on the command line:  --parallelism.data_parallel_shard_degree 2 --parallelism.tensor_parallel_degree 2

And here is the core of FP8, from DeepSeek-V3's inference kernel. Each tile is scaled so its largest value maps to 448, the largest number E4M3 can hold.

@triton.jit
def act_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr, scale_fmt: tl.constexpr):
    pid = tl.program_id(axis=0)
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)      # one 1x128 tile of activations
    x = tl.load(x_ptr + offs).to(tl.float32)
    amax = tl.max(tl.abs(x))
    amax = tl.maximum(amax, 1e-4)
    s = amax / 448.                                          # 448 = the largest e4m3 value
    y = x / s
    tl.store(y_ptr + offs, y.to(y_ptr.dtype.element_ty))     # float8_e4m3fn
    tl.store(s_ptr + pid, s)                                 # one fp32 scale per tile

The memory arithmetic

What the communication looks like

• FSDP: one all-gather per layer forward, one all-gather and one reduce-scatter per layer backward, overlapped with the next layer's compute. reshard_after_forward frees the gathered weights between forward and backward; torchtitan keeps them for the last block and under pipeline parallel.
• TP: two all-reduces per block (after attention, after the MLP), or reduce-scatter and all-gather pairs with sequence parallel. Needs NVLink; never crosses a node. Async TP overlaps these with the matmuls.
• PP: activations pass directly from each stage to the next. The "bubble" is the fraction of time stages sit idle: (stages − 1) / micro-batches for 1F1B. That is why runs use many micro-batches, and why zero-bubble and DualPipe schedules exist.
• EP: two all-to-alls per MoE layer, dispatch and combine. DeepEP is DeepSeek's kernel for it, with FP8 dispatch; node-limited routing caps how many nodes a token's experts span.
• CP: keys and values pass around a ring. Under causal attention, later tokens do more work than early ones, so torchtitan gives each GPU one chunk from the start of the sequence and one from the end to even out the load.

Where it breaks

• Order of wrapping matters. torchtitan applies declarative sharding (TP), then activation checkpointing, then compile, then FSDP. Get it wrong and compile graphs break or hooks fire twice.
• Clipping needs a global norm. The stock clip_grad_norm_ only sees one rank's slice under PP; torchtitan sums norm^p across stages and then takes the p-th root. nanochat once clipped each rank separately instead of globally, and its log records the bug.
• The router must not be quantized. torchtitan's float8 converter refuses to touch the router gate, and the 671B config leaves lm_head and router.gate out of FP8. Small matmuls also get slower in FP8, because the cost of quantizing outweighs the gain, so there is an automatic filter by shape.
• FP8 needs compile. Without it, the scale-and-cast steps run as separate kernels and FP8 ends up slower than bf16. Rowwise scaling keeps all communication in high precision; only the matmuls run in e4m3.
• MFU is what you tune for. MFU is the share of the GPUs' peak speed you are using. Picotron reports 38 percent on Llama 2 7B across 64 H100s and about 50 percent on a 1.7B model on 8. Anything above 40 percent at scale is good; MoE and long-context runs come in lower.

Checkpoints

Every 500 to 1,000 steps, save the model, the optimizer state and the dataloader position to storage. PyTorch's distributed checkpoint (DCP) has every GPU write its own shard in parallel. It can also load the result onto a different number of GPUs, which lets you resume a 512-GPU run on 480 after a node dies. Saving happens in the background. The weights are first copied from GPU memory to host memory (staging), a background process writes them out, and the next optimizer step waits only for the copy to finish. OLMo 3 saved every 1,000 steps. Each stage of its run starts from the URL of the previous stage's checkpoint, and every intermediate checkpoint is published.

Loss spikes

In a loss spike, the loss suddenly jumps and sometimes never comes back down. The cause is almost always attention logits growing until the softmax saturates. There are three defenses, in order of preference. First, prevent spikes in the architecture: QK-norm, reordered norm, z-loss, no weight decay on embeddings. Second, prevent them in the optimizer: MuonClip, which is how Kimi K2 got through 15.5T tokens with zero spikes. Third, skip the bad step: OLMo-core's SkipStepOptimizer skips any update whose loss or gradient norm is more than 6 standard deviations above the mean of the last 128 steps. DeepSeek-V3 reports no spikes it could not recover from and no rollbacks. torchtitan stops on a NaN or infinite loss and makes you restart from a checkpoint.

What to watch

Watch the loss, smoothed. Watch the gradient norm, because it rises before a spike. Watch tokens per second and MFU (model FLOPs utilization): the fraction of the hardware's peak speed you are using, computed from a formula rather than measured. Check held-out loss every few hundred steps, in bits per byte. And run downstream evals every few thousand steps: HellaSwag, ARC, MMLU subsets, a code task and a math task. Each is scored by how likely the model finds the right continuation, because a base model cannot follow instructions.

CORE

nanochat defines "GPT-2 grade" by the DCLM CORE score. CORE has 22 few-shot tasks. Each task is scored as centered accuracy, (accuracy − random baseline) / (1 − random baseline), and the task scores are averaged. GPT-2 1.6B scores 0.2565, and the current speedrun beats it in 1.65 hours. A single number that moves smoothly over a run is more useful than twenty numbers that jump around.

DCP, from torchtitan's checkpointer. The to_hf branch writes safetensors that Hugging Face tooling can load.

if to_hf:
    state_dict = self.sd_adapter.to_hf(state_dict)                    # rename keys into the HF layout
    storage_writer = HuggingFaceStorageWriter(path=save_path, save_distributed=True, ...)
if async_mode == AsyncMode.ASYNC:
    ret = dcp.async_save(state_dict, storage_writer=storage_writer,
                         checkpoint_id=checkpoint_save_id, process_group=self.pg)
elif async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:                     # stage to pinned host memory first
    ret = dcp.async_save(state_dict, storage_writer=storage_writer, checkpoint_id=checkpoint_save_id,
                         process_group=self.pg, async_checkpointer_type=AsyncCheckpointerType.PROCESS,
                         async_stager=self.stager)
else:
    ret = dcp.save(state_dict, storage_writer=storage_writer, checkpoint_id=checkpoint_save_id)

nanochat saves less, but saves the thing people forget: where the dataloader was.

checkpoint = {
    "model": orig_model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "dataloader": dataloader_state_dict,            # {"pq_idx": ..., "rg_idx": ..., "epoch": ...}
    "loop": {"step": step, "min_val_bpb": min_val_bpb, "smooth_train_loss": smooth_train_loss,
             "total_training_time": total_training_time},
}

MFU is simple arithmetic. Count six FLOPs per parameter per token for the matmuls (forward plus backward), add the attention term, multiply by tokens per second, and divide by the peak speed the vendor publishes for dense bf16.

# nanochat GPT.estimate_flops, generalized for sliding windows
num_flops_per_token = 6 * self.num_matmul_params()                    # embeddings excluded
for layer in layers:
    num_flops_per_token += 12 * n_head * head_dim * min(window, T)      # attention: q·k and softmax·v
mfu = 100 * (num_flops_per_token * tokens_per_sec) / (peak_flops * world_size)
# peak_flops: H100 SXM 989e12 · H100 PCIe 756e12 · A100 312e12   (dense bf16, no sparsity)

And here is the CORE calculation, simplified.

def core_score(results):                                   # nanochat/core_eval.py, the idea
    centered = []
    for task, acc in results.items():
        baseline = RANDOM_BASELINE[task]                    # e.g. 0.25 for a 4-way multiple choice
        centered.append((acc - baseline) / (1 - baseline))  # 0 = random, 1 = perfect
    return sum(centered) / len(centered)                    # mean over 22 tasks

What a month on a cluster is like

• Failures happen all the time. Llama 3's paper counted 466 job interruptions over 54 days of the 405B run, 419 of them unexpected, most from GPU faults. The training loop is a small part of the work. The rest is the automation that notices a dead node, restarts from the last checkpoint and gets back to full speed.
• Checkpoints cost time, so budget them. Pick the interval so that the work you expect to lose to failures (half an interval, times the failure rate) balances the time spent saving. Async staging makes a save nearly free in wall-clock time; the storage bandwidth per GPU sets the lower limit.
• A seed checkpoint makes runs reproducible. torchtitan writes the initialized weights once so that a run can be restarted from identical numbers on any GPU count.
• Resume is approximate unless you make it exact. nanochat advances one row group past the saved position to avoid repeating data; OLMo's index makes the position exact. The random number generator state is rarely saved, so having no dropout helps here too.
• Metrics can mislead. DCLM changed CORE from v1 to v2 in September 2025, so scores from before and after that change cannot be compared. Held-out loss depends on the tokenizer, which is why people use bits per byte. Benchmark questions that leak into the training data inflate every score.
• Souping (averaging checkpoints) is normal. OLMo 3 32B averages the weights of two midtraining runs, and of its last three long-context checkpoints. The published "base" is often an average of several checkpoints rather than a single step.

A base model does completion. Given "The capital of France is", a good one writes " Paris." Then it keeps going, because it is still predicting the most likely continuation: " The capital of Germany is Berlin." Given a question formatted like a forum post, it writes a forum answer, then another post. The only thing that makes it stop is the document separator, which it learned to predict at the end of documents.

Before post-training existed, people used few-shot prompting: show the model three examples of the format you want, and it completes a fourth. That is still how base models are evaluated. HellaSwag, MMLU and the CORE tasks are all written out as a few examples plus candidate continuations, and scored by the probability the model gives the right one. A base model can score well on knowledge benchmarks and still be useless in a chat.

nanochat's README describes the speedrun model as "a bit like talking to a kindergartener": it knows the sky is blue and will make up the reason. That is what 1.65 hours of H100 time buys. The frontier base models are the same kind of model, trained a few million times longer.

Sampling from a base model is the inference loop with nothing around it: one forward pass per token, take the argmax or sample, append, repeat.

ids = [tokenizer.get_bos_token_id()] + tokenizer.encode("The capital of France is")
for _ in range(16):
    logits = model(torch.tensor([ids]))[0, -1]                  # only the last position matters
    next_id = int(logits.argmax())                               # temperature 0: greedy
    ids.append(next_id)
    if next_id == tokenizer.get_bos_token_id():                  # the document separator: "this text is over"
        break
print(tokenizer.decode(ids[1:]))

The evaluation harness renders benchmarks the same way. A multiple-choice question becomes several candidate strings, and the score is which candidate the model finds most probable.

context = few_shot_examples + question                          # 5 solved examples, then the real one
scores = []
for choice in choices:                                          # "A) ...", "B) ...", ...
    ids = tokenizer.encode(context + choice)
    logp = model.log_prob_of_continuation(ids, start=len(tokenizer.encode(context)))
    scores.append(logp)                                         # sum of log-probs of the choice tokens
correct = scores.index(max(scores)) == answer_idx

The handoff

Which checkpoint becomes "the base model" is a choice. Usually it is the one from after the anneal, when the learning rate has decayed to zero on the high-quality mix. Sometimes it is an average of several. OLMo 3 publishes every stage: Base, then SFT, then DPO, then RLVR, for both an Instruct track and a Think track, and every intermediate checkpoint can be downloaded. Kimi K2's report recommends fine-tuning a checkpoint pretrained with Muon using Muon again. Kimi K3 does quantization-aware training and tunes its multi-token-prediction draft head during post-training, so the served model computes with the same number formats it was trained with.

The base checkpoint goes into supervised fine-tuning first, where it learns the chat format and the special tokens the tokenizer reserved back in step 2. That is the next page.