Memory: Kimi K3
on 512 GB200s

FP32 training does not mean FP32 GEMMs

In this recipe, the sharded resident parameter is the master weight. FSDP makes a BF16 compute view for forward and backward; the optimizer updates the FP32 resident shard. Do not count a separate, full-model BF16 master copy on every GPU.

Muon needs one momentum, not two Adam moments

The first study counted 2,779,931,738,208 parameters. In this optimizer assignment, 2,777,072,664,576 use DistMuon and 2,859,073,632 use AdamW: about 99.897% versus 0.103%. Routed expert matrices alone are 97.94% of the model. Selecting 16 of 896 experts reduces compute per token, not the number of trainable weights or optimizer states we must store.

Update-state tally = 12 P_muon + 16 P_adam
                   = 33,370,617,153,024 bytes
                   = 30.35 TiB across the logical model

Ideal 512-way average = 60.70 GiB / GPU

This tally assumes all trainable parameters have gradients and optimizer state. It is not a measured peak or an assertion that every gradient exists throughout forward. It excludes BF16 compute views, unsharded accumulation, activations, padding, communication, and optimizer workspaces. The familiar “16 or 17 bytes per parameter” shortcut is not a universal memory law.

A memory lab, with its assumptions exposed

Pick a pipeline rank. The state ledger uses the actual 93-layer structure, including the first dense FFN, last MLA layer, embeddings, vision pathway, and output head. The backward budget adds explicitly modeled working sets. It is a planning model, not a substitute for measuring an implementation.

PP, EP, and FSDP divide different things

On one pipeline stage, let D = non-routed parameters and E = routed parameters.

Resident parameters / GPU = D / dense_FSDP + E / (EP * expert_FSDP)
BF16 gathered stage view  = 2 * (D + E / EP) bytes

At PP8 / EP64 / FSDP64: expert_FSDP = 1

Increasing EP at fixed PP and GPU count can shrink the expert compute slice while reducing the expert FSDP degree. It does not create additional aggregate HBM. And the 512 GPUs do not form one pooled allocation space: every rank must stay within its own device capacity.

64 microbatches do not mean 64 saved forwards

In ordinary 1F1B, a rank fills its local pipeline, then alternates forward and backward. Backward releases that microbatch's saved activations. With eight physical stages and enough microbatches, rank 0 can hold eight unfinished forwards; rank 7 needs only one. This is a logical schedule bound, excluding implementation-specific send/receive buffers and keepalives.

Peak unfinished forwards on rank r = min(microbatches, PP - r)

At the default shape, local packed tokens T = rows × context = 4096. One BF16 [T, 7168] stream is 56 MiB. Twelve checkpointed layer inputs across five unfinished forwards on rank 3 give a boundary-input floor of 3.28 GiB. The growing AttnRes bank and kernel workspaces are additional, not hidden inside that number.

Global tokens per update are DP × microbatches × rows × context: 64 × 64 × 1 × 4096 = 16,777,216. Increasing microbatches beyond the 1F1B live-forward bound mostly extends the schedule; increasing rows or context increases each microbatch's activation and routing footprint. Two 2K rows and one 4K row have the same flattened token count, but different document attention boundaries.

KDA's fixed-size inference state does not make its training activations constant-size: token and chunk intermediates still need storage or recomputation. Varlen FlashAttention avoids materializing the full attention-probability matrix for MLA, but still needs Q/K/V, outputs, and backward statistics. Neither attention path has a zero activation-memory bill.

From a 30-TiB state bill to a local budget

Each step changes one ownership or activation assumption. The same FP32-master / FP32-state / BF16-compute policy is kept throughout. Component lengths below use a logarithmic axis; add the numeric values, not bar lengths.

A measured run is the capacity proof

The reference completed 12 updates with FP32 resident state, BF16 compute, PP8 / FSDP64 / EP64, full layer checkpointing, varlen attention, compiled pointwise regions, no CUDA graphs, and forced-balanced HybridEP routing. It demonstrates short-run feasibility of this specific configuration; it is not a general capacity guarantee or a convergence result. The fixed reference stays fixed when you change calculator controls.

The calculator deliberately does not label the gap between its subtotal and 146.18 GiB as “activation memory.” That would invent a decomposition. Some important reasons implementations exceed a state-only estimate:

Tensor allocation, allocator reservation, and total device usage are different measurements. The reported 38.13-GiB difference ignores external allocations and fragmentation. Check every PP rank and the first optimizer update, not only the first forward. Also check later iterations, when momentum and caches are populated.