A. saved-for-backward bytes (float32)
  L  B    T    d  h    measured  checkpointed   ratio
  2  2   64   32  4      815104         49152   16.58
  2  2  128   32  4     2170880         98304   22.08
  2  4   64   64  4     2703360        196608   13.75
  4  2  128   64  4     6438912        327680   19.65
  4  2  256   64  8    25591808        655360   39.05
  4  1  256  128  8    17186816        655360   26.23
  6  2  128   96  6    14376960        688128   20.89
  2  2  512   32  4    21659648        393216   55.08
fit: bytes = 16.13 * (L B T d * 4) + 1.03 * (L B h T^2 * 4);  max relative error 0.020

B. KV cache formula  bytes = 2 * layers * kv_heads * d_head * tokens * bytes_per_element
  allocated bytes match the formula

C. hypothetical 5.9B configuration (bf16 weights)
  parameters 5.90 B; bf16 weights 11.0 GiB
  AdamW training state, mixed precision (bf16 weights 2 + bf16 grads 2 + fp32 master 4 + fp32 m, v 8 = 16 bytes/param): 87.9 GiB; sharded over 8 devices: 11.0 GiB each
  KV cache per token 128 KiB; per 4096-token sequence 0.50 GiB; without grouped-query sharing (32 kv heads) 2.00 GiB
  sequences of 4096 tokens that fit on an 80 GiB device after weights and 6 GiB of workspace: 126 (grouped-query), 31 (32 kv heads)
  training activations per token per layer (measured coefficient a = 16.1 units of d x 4 bytes in float32, i.e. 16.1 d elements): 129 KiB in bf16 (attention-score term b = 1.03 h T^2 elements per layer: 264 KiB per token at T = 4096 if scores are materialized)
  per-token activations, no attention scores: full 4.03 MiB; with per-block checkpointing 0.38 MiB (ratio 10.7); 8192 tokens per device: 32.3 GiB vs 3.0 GiB
