File size: 39,388 Bytes
439e7b1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 | # model-math
## SUMMARY
Config A (24 layers, d=1024, 16 heads, FFN 4096, V=32768, Lmax=512, pre-LN, tied MLM head) is the recommended pin: 337,476,608 params for the pretraining model (302.3M encoder + 34.1M embeddings + 1.08M MLM head), 337,443,842 for the deployed classifier; it is +2.3% over the 330M target, and the nearest-to-330M alternatives are 23 layers (324.9M) or V=16384 (320.7M). Training state is 16 B/param = 5.03 GiB (+0.63 GiB autocast bf16 weight cache); with SDPA flash/mem-efficient attention the per-token-per-layer activation footprint is about 38·d bytes under torch.autocast (34·d if the residual stream is bf16), versus Korthikanti's 34·d + 5·a·s for eager attention, so on the 24 GiB card with 1.9 GiB of desktop use and a 4 GiB headroom target the no-checkpoint micro-batch is 16 (24 is marginal), while full per-layer checkpointing fits 48+ with 4 GiB to spare; computing MLM logits only on the ~15% masked positions saves 85% of the logits/cross-entropy memory (136 MiB per sample at B=1). The RTX 3090 Ti bf16 tensor-core peak is 80 TFLOPS dense with FP32 accumulate (160 sparse; 40 TFLOPS is the non-tensor rate) — derived from the GA102 whitepaper's RTX 3090 figure of 71/142 scaled by SM count and clock, and consistent with Wikipedia's 320 sparse FP16-accumulate figure halved for FP32 accumulate; at 2.03 GFLOP/token (6N + 14·nl·d·L + masked head) and 30-45% MFU that is 11.8k-17.8k tokens/s, i.e. 2.3-3.5 h per epoch of 150M tokens (1,144 optimizer steps at B_eff=256) and 6.3-9.4 h for 400M, with full checkpointing costing x1.33. Batch-1 inference of a 512-token email is ~336 GFLOP: ~14 ms on the 3090 Ti at 30% MFU but ~336 ms on the 5950X at 1 TFLOPS effective, so the 74 ms/sample benchmark is only met on a GPU (ZeroGPU runs on RTX Pro 6000 Blackwell slices) or with truncation to ~128 tokens on CPU; real emails averaging 200-300 tokens make dynamic padding worth roughly 2x. Recommended training recipe: checkpointing off by default, micro-batch 16 x 16 accumulation = 256, AdamW peak LR 1.5e-4 (range 1e-4 to 2e-4), warmup min(1000 steps, 6% of total), betas (0.9, 0.98), eps 1e-6, weight decay 0.01, clip 1.0, 15% dynamic masking 80/10/10, sequence packing for pretraining (also keeps the Flash SDPA backend eligible because it needs no attn_mask) and length-bucketed dynamic padding for fine-tuning and inference.
## FINDINGS
# Architecture / parameter / VRAM / throughput math for the from-scratch encoder
Calculator: `/tmp/claude-1000/-home-nathank-BITS-dev-Phishing-Model/8697d948-f234-4e8c-a1e6-c0b6f6ddd7ef/scratchpad/model_math.py` (pure Python 3.14.7, no torch; full output in `model_math.out` next to it and pasted in section 7). Nothing was installed and nothing under the project directory was touched.
## 0. Hardware facts verified live (needed before any throughput number)
- **RTX 3090 Ti**: 84 SMs, 10,752 CUDA cores, 336 tensor cores, boost 1860 MHz, FP32 39.99 TFLOPS, "Tensor compute (FP16)" 320 TFLOPS *sparse* (table formula: tensor cores x clock x 256/1000 = FP16-accumulate), 24 GB @ 1008 GB/s, 450 W — https://en.wikipedia.org/wiki/GeForce_RTX_30_series
- **GA102 whitepaper Table 9 (RTX 3090 FE, 82 SMs @ 1695 MHz)**: "Peak FP16 Tensor TFLOPS with FP16 Accumulate 142/284", "Peak FP16 Tensor TFLOPS with FP32 Accumulate 71/142", "Peak BF16 Tensor TFLOPS with FP32 Accumulate 71/142" (second number = sparsity). The RTX A6000 rows (same GA102, 84 SMs) show 154.8/309.6 for *both* FP16- and FP32-accumulate, i.e. **GeForce GA10x runs FP32-accumulate tensor math at half rate; pro cards do not**. No prose explains it; the table rows are the evidence — https://www.nvidia.com/content/PDF/nvidia-ampere-ga-102-gpu-architecture-whitepaper-v2.pdf (Appendix A, p.44)
- **Derived 3090 Ti bf16 dense peak**: 71 x (84x1860)/(82x1695) = **79.8 ≈ 80 TFLOPS** (160 sparse). Cross-check: 336 x 1.860 x 256 = 160 TFLOPS FP16-accumulate dense, halved for FP32 accumulate = 80. **Answer to "80 or 40?": 80 TFLOPS dense bf16 (FP32 acc); 40 TFLOPS is the non-tensor FP32/BF16 rate.** There is no bf16-with-bf16-accumulate mode in the table, so bf16 GEMMs are always the 80 figure.
- PyTorch bf16 GEMMs: `torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction` "is set to True by default for BF16" (affects split-K reductions only; tensor-core accumulation is fp32) — https://docs.pytorch.org/docs/2.14/notes/cuda.html. Same page: `memory_allocated` vs `memory_reserved`, and `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` to reduce fragmentation when batch shapes change.
- TechPowerUp returned HTTP 403 (not used). Everything attributed to Context7 was cross-checked against docs.pytorch.org 2.14 pages (upstream wins).
## 1. Exact parameter counts
Formulas (stated assumptions): token emb V·d; position emb Lmax·d; emb LN 2d; per layer QKV 3d²+3d, out-proj d²+d, 2 LN 4d, FFN up d·f+f, FFN down f·d+d; final LN 2d (pre-LN needs it); MLM head dense d²+d + LN 2d + output bias V (tied) [untied: +V·d]; classifier d²+d + 2d+2; no segment embeddings.
Sanity check: applying the same per-layer/embedding formulas to BERT-base (post-LN, V=30522, +2 segment rows, +pooler) gives exactly **109,482,240**, which matches `BertModel('bert-base-uncased').num_parameters()` (HF figure from memory, UNVERIFIED live; the exact match is strong evidence the formulas are right). BERT paper quotes 110M/340M — https://ar5iv.labs.arxiv.org/html/1810.04805
| config | per-layer | encoder (+final LN) | embeddings | MLM head tied | **TOTAL tied** | untied | fine-tune (cls head) | N_nonemb |
|---|---|---|---|---|---|---|---|---|
| **A 24L d1024 h16 f4096 V32768** | 12,596,224 | 302,311,424 | 34,080,768 | 1,084,416 | **337,476,608** | 371.0M | **337,443,842** | 303.4M |
| B 24L V32000 | 12,596,224 | 302,311,424 | 33,294,336 | 1,083,648 | 336,689,408 | 369.5M | 336.7M | 303.4M |
| C 24L V16384 | 12,596,224 | 302,311,424 | 17,303,552 | 1,068,032 | 320,683,008 | 337.5M | 320.7M | 303.4M |
| D 20L V32768 | 12,596,224 | 251,926,528 | 34,080,768 | 1,084,416 | 287,091,712 | 320.6M | 287.1M | 253.0M |
| E 22L V32768 | 12,596,224 | 277,118,976 | 34,080,768 | 1,084,416 | 312,284,160 | 345.8M | 312.3M | 278.2M |
| F 24L d768 h12 f3072 V32768 | 7,087,872 | 170,110,464 | 25,560,576 | 624,896 | 196,295,936 | 221.5M | 196.3M | 170.7M |
| A' 23L d1024 V32768 (extra) | 12,596,224 | 289,715,200 | 34,080,768 | 1,084,416 | 324,880,384 | 358.4M | 324.8M | 290.8M |
Config A breakdown: layer = qkv 3,148,800 + out_proj 1,049,600 + 2xLN 4,096 + ffn_up 4,198,400 + ffn_down 4,195,328; embeddings = tok 33,554,432 + pos 524,288 + LN 2,048; MLM head tied = dense 1,049,600 + LN 2,048 + bias 32,768; classifier head = 1,051,650.
Distance from 330M (tied total): A' 23L −5.12M; B +6.69M; **A +7.48M**; C −9.32M; E −17.72M; D −42.91M; F −133.70M. **Recommend A** (BERT-large shape, 2^15 vocab, +2.3% over target; BERT-large itself is 335–340M). If the brief enforces ≤330M strictly, use A' (23L, 324.9M) rather than C (16K vocab fragments URLs/headers badly). Deployed classifier = 337,443,842 params → 675 MB bf16 / 1.35 GB fp32 safetensors.
Vocab alignment: 32768 = 2^15 (multiple of 64 and 128), 32000 = 64·500 = 128·250, 16384 = 2^14 — all fine. NVIDIA guidance: for FP16/BF16 "matrix dimensions would need to be multiples of 8 elements for best efficiency (or 64 elements on A100)", and "choosing multiples of larger powers of two up to 128 bytes (64 for FP16) can further improve efficiency"; tile and wave quantization explain the cliffs — https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html
## 2. Memory
### Static (AdamW, bf16 autocast, fp32 master params)
16 B/param (fp32 params 4 + fp32 grads 4 + Adam m,v 8) x 337,476,608 = **5.03 GiB (5.40 GB)**; plus autocast's weight-cast cache (`cache_enabled`, bf16 copies of parameters inside the autocast region) ≈ 2 B/param = 0.63 GiB → **5.66 GiB**. Add CUDA context 0.5–1 GiB and the 1.9 GiB already used by kwin/plasma. GPU = 24,564 MiB = 23.99 GiB. Budget left for activations+logits with a 4 GiB headroom target: **11.68 GiB** (process cap 18.09 GiB).
### Per-token-per-layer activation bytes saved for backward (pre-LN, autocast bf16, SDPA)
Korthikanti et al. 2022 Eq.1: "Activations memory per layer = sbh(34 + 5as/h)" for fp16 eager attention; itemized: attention 11sbh + 5as²b (QKV input 2sbh, Q,K 4sbh, softmax out 2as²b, dropout mask as²b, dropout out 2as²b, V 2sbh, attn-over-V out 2sbh, out dropout mask sbh), MLP 19sbh (input 2, GELU input 8, GELU output 8, mask 1), LNs 4sbh — https://ar5iv.labs.arxiv.org/html/2205.05198. Per token per layer that is 34d + 5·a·s = 34,816 + 40,960 = 75,776 B at d=1024, a=16, s=512 — the L×L terms would be **+118% on top of 34d**.
FlashAttention stores only "the output O and the softmax normalization statistics (m,ℓ)" and recomputes S and P in backward; dropout keeps "the pseudo-random number generator state" instead of a mask; memory is linear in sequence length — https://ar5iv.labs.arxiv.org/html/2205.14135. PyTorch SDPA dispatches to FlashAttention-2 / Memory-Efficient / C++ math — https://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html. FlashAttention-2 supports "Ampere, Ada, or Hopper GPUs (e.g., A100, RTX 3090, RTX 4090, H100)", "fp16 and bf16", "all head dimensions up to 256" — https://github.com/Dao-AILab/flash-attention (3090 Ti is sm86 Ampere, head_dim 64: eligible).
My tally (config A), tensors actually saved:
| tensor | fp32 residual (autocast) | bf16 residual |
|---|---|---|
| LN1 input = residual stream | 4d | 2d |
| QKV linear input (bf16 LN1 out) | 2d | 2d |
| Q,K,V bf16 (saved by SDPA) | 6d | 6d |
| attn output O bf16 (SDPA bwd + out-proj input, one tensor) | 2d | 2d |
| SDPA logsumexp fp32 | 4h (64 B) | 4h |
| out-proj dropout mask (bool) | d | d |
| LN2 input = residual stream | 4d | 2d |
| FFN-up input (bf16 LN2 out) | 2d | 2d |
| FFN-up out = GELU in (bf16, width 4d) | 8d | 8d |
| GELU out = FFN-down in (bf16) | 8d | 8d |
| FFN-down dropout mask | d | d |
| **total** | **38,976 B = 38.06 d (c = 38 + 4h/d)** | **34,880 B = 34.06 d (c = 34)** |
Why 38 not 34: under `torch.autocast` `layer_norm`, `softmax`, `log_softmax`, `cross_entropy`, `nll_loss` are on the "autocast to float32" list while `linear`/`matmul` are on the float16/bf16 list — https://docs.pytorch.org/docs/2.14/amp.html. Embedding lookups and residual adds are not autocast ops, so the residual stream stays fp32 and each LN saves a 4-byte input (+2d per LN). Use **c = 38** for planning; c = 34 only if the model is run with a bf16 residual stream. An eager softmax under autocast runs in fp32, so an accidental math-backend fallback costs ≥7·a·s per token per layer (worse than the paper's 5·a·s).
Non-layer per token: 9d (fp32 emb output/emb-LN input 4d, emb dropout mask d, final-LN input 4d). MLM head per *scored* position: 12d + logits 2V (bf16) + fp32 cross-entropy intermediates (4V saved log_softmax output + 4V transient fp32 up-cast) = **10V B/position at peak**. Per 512-token sample: no-ckpt layers 457 MiB (c=38) / 409 MiB (c=34); full-ckpt fp32 layer inputs 48 MiB (+19 MiB one live layer) / bf16 inputs 24 MiB; logits full 160 MiB vs masked ~15% (76 pos) 24 MiB → **gather-before-head saves 136 MiB/sample (85%)**.
### Fit tables (config A, L=512; totals include desktop 1.9 + ctx 0.75 + static 5.66 + 10% allocator slack on activations; FITS = ≥4 GiB headroom, tight = 2–4, NO < 2)
(i) no checkpointing, c=38d (autocast): B=4 masked 10.39 GiB total (13.60 headroom) FITS; B=8 12.47 (11.52) FITS; **B=16 masked 16.63 (7.36) FITS**, B=16 full-logits 18.84 (5.15) FITS; **B=24 masked 20.78 (3.20) tight**, B=24 full 24.10 NO; B=32 masked 24.94 NO; B=48 NO.
(i') c=34d (bf16 residual): B=16 masked 15.80 (8.19) FITS; B=24 masked 19.55 (4.44) FITS; B=32 masked 23.29 (0.69) NO.
(ii) full per-layer checkpointing, fp32 layer inputs 4d/tok/layer + one live layer: B=16 masked 9.93 (14.06); B=32 masked 11.55 (12.44); **B=48 masked 13.17 (10.82) FITS**, B=48 full-logits 19.81 (4.18) FITS.
(ii') checkpointing with bf16 layer inputs (2d, the task's assumption): B=48 masked 11.83 (12.16) FITS; B=48 full 18.46 (5.52) FITS.
**Conclusion**: within a 14–18 GiB process budget and ≥4 GiB headroom, (i) no-checkpoint → **micro-batch 16** (24 only if the residual is bf16 or the desktop frees memory); (ii) full checkpointing → **48 fits with >10 GiB spare, 64–96 would too**. Checkpointing is not needed for memory at B=16 and costs x1.33 compute, so default it OFF.
## 3. Training throughput
FLOPs/token: Kaplan: N "excluding all vocabulary and positional embeddings", "C_forward ≈ 2N + 2 n_layer n_ctx d_model", "C ≈ 6NBS" — https://ar5iv.labs.arxiv.org/html/2001.08361. Chinchilla: "FLOPs(N,D) ≈ 6ND", Appendix F exact count within 0.99–1.10 of 6ND — https://ar5iv.labs.arxiv.org/html/2203.15556. PaLM App. B: "6N matmul FLOPs per token seen (2N for forward and 4N for backward)" and "dense self-attention add 6LH(2QT) FLOPs per token", MFU = tokens/s ÷ P/(6N + 12LHQT) — https://arxiv.org/pdf/2204.02311. With flash recompute of QK^T in backward the attention term is 14·nl·d·L (12 standard).
Config A, L=512, N_nonemb = 303.4M: 6N = 1.820 GF; attention 14·24·1024·512 = 0.176 GF (12·nl·d·L = 0.151 GF; 8.7% of total); tied decoder on 15% masked positions 6·V·d·0.15 = 0.030 GF (0.201 GF if scored on all 512) → **2.027 GFLOP/token** (2.198 full logits). Full checkpointing: 2.694 GFLOP/token = **x1.329**.
Peak 80 TFLOPS; MFU 30 / 37.5 / 45% → 24 / 30 / 36 TFLOPS → **11,843 / 14,804 / 17,764 tokens/s**. B_eff = 256 x 512 = 131,072 tokens/step.
| corpus | steps/epoch | h/epoch @30% | @37.5% | @45% | ckpt @30% | ckpt @45% |
|---|---|---|---|---|---|---|
| 100M | 763 | 2.35 | 1.88 | 1.56 | 3.12 | 2.08 |
| 150M | 1,144 | 3.52 | 2.81 | 2.35 | 4.68 | 3.12 |
| 250M | 1,907 | 5.86 | 4.69 | 3.91 | 7.79 | 5.20 |
| 400M | 3,052 | 9.38 | 7.51 | 6.25 | 12.47 | 8.31 |
A 1–2k-step warmup exceeds one epoch for corpora ≤250M at B_eff=256 → plan several epochs (BERT: 1M steps of 256x512 ≈ 40 epochs over 3.3B words — https://ar5iv.labs.arxiv.org/html/1810.04805) or use B_eff=128 to double the step count. MFU 30–45% is a planning assumption for eager PyTorch + SDPA without torch.compile; measure with `torch.cuda.Event` timing on the first 50 steps.
## 4. Inference latency (fine-tuned classifier, batch 1)
FLOPs = 2·(encoder+cls)·L + 4·nl·d·L²: L=512 → 336.4 GFLOP; 256 → 161.8; 128 → 79.3.
- 3090 Ti bf16 @30% MFU (24 TFLOPS): **14.0 / 6.7 / 3.3 ms** (plus a kernel-launch floor of a few ms for 24 layers at batch 1 — UNVERIFIED, measure).
- 5950X fp32 @1.0 TFLOPS effective: **336 / 162 / 79 ms**; @0.5 TFLOPS: 673 / 324 / 159 ms. (Theoretical 16 cores x 32 FLOP/cycle x ~3.9 GHz ≈ 2.0 TFLOPS; 1 TFLOPS effective ≈ 50%, UNVERIFIED assumption from the brief.)
- vs instructor benchmark **74 ms/sample**: GPU beats it ~5x at 512 tokens; CPU misses it by ~4.5x at 512 and ~2x at 256, only borderline at 128 tokens. Real emails averaging 200–300 tokens (task-given, UNVERIFIED) halve dense cost and quarter attention cost → dynamic truncation/padding to the longest in batch (rounded to a multiple of 8) matters.
- ZeroGPU: default `large` = "Half NVIDIA RTX Pro 6000 Blackwell, 48GB" (xlarge = full, 96GB, 2x quota); PRO can host 10 ZeroGPU Spaces, 40 min/day quota, `@spaces.GPU` default 60 s duration, Gradio SDK only, PyTorch 2.8.0–2.13.0 and Python 3.12.12 / 3.10.13 supported, `torch.compile` not supported (AOT compile only), models must be moved to `cuda` at module level — https://huggingface.co/docs/hub/spaces-zerogpu. Batch-1 latency there will be dominated by queue/cold start, not FLOPs.
## 5. Recommendations to pin
- **Model**: n_layers=24, d=1024, heads=16 (head_dim 64), ffn=4096, vocab=32768, max_len=512, pre-LN, GELU, dropout 0.1, tied MLM output, no segment embeddings → 337,476,608 params (pretrain) / 337,443,842 (classifier).
- **Gradient checkpointing**: default OFF (B=16 fits with 7.4 GiB headroom; checkpointing costs x1.33). Expose a flag; with it ON start the batch finder at 48–64. Use `torch.utils.checkpoint.checkpoint(..., use_reentrant=False)` ("Recommended to use False"; early_stop default) — https://docs.pytorch.org/docs/2.13/checkpoint.html.
- **Micro-batch / auto batch finder**: start at 16 (x16 accumulation = 256), probe 24 (expect tight), fall back 12 → 8. Verify with `torch.cuda.max_memory_reserved()` after 3 steps; set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`.
- **Effective batch**: 256 sequences x 512 = 131,072 tokens/step (BERT's exact setting).
- **Optimizer / LR**: AdamW, betas (0.9, 0.98), eps 1e-6 — RoBERTa: "setting β₂=0.98 to improve stability when training with large batch sizes", ε tuned, Table 9 LARGE peak LR 4e-4 at batch 8K, 30K warmup, wd 0.01, dropout 0.1 — https://ar5iv.labs.arxiv.org/html/1907.11692; BERT: LR 1e-4 at batch 256, warmup 10k, linear decay, wd 0.01 — https://ar5iv.labs.arxiv.org/html/1810.04805. Pre-LN tolerates larger LR and "the learning-rate warm-up stage can be safely removed"; their pre-LN BERT-base used 3e-4 without warmup — https://ar5iv.labs.arxiv.org/html/2002.04745. For a 24-layer pre-LN model at batch 256 from scratch: **peak 1.5e-4 (range 1e-4–2e-4), warmup min(1000 steps, 6% of total), linear or cosine decay to 10% of peak, wd 0.01 excluding biases/LN, grad-clip 1.0**. Fine-tune: LR 2e-5–5e-5, batch 16–32, 2–4 epochs (BERT ranges).
- **MLM**: 15% dynamic masking per batch (RoBERTa), 80% [MASK] / 10% random / 10% unchanged (BERT), loss only on masked positions, gather hidden states at masked positions before the MLM head (85% head memory saving, 0.17 GFLOP/token saving).
- **Sequence packing**: YES for pretraining — RoBERTa FULL-SENTENCES (contiguous text crossing document boundaries, no NSP; DOC-SENTENCES was slightly better but FULL-SENTENCES chosen for simplicity). Packing also removes padding and the attn_mask, which keeps the Flash backend eligible: PyTorch's dispatcher warns "Flash Attention does not support non-null attn_mask." — https://raw.githubusercontent.com/pytorch/pytorch/main/aten/src/ATen/native/transformers/sdp_utils_cpp.h; with a boolean key-padding mask SDPA uses the memory-efficient backend (still no L×L storage). NO packing for fine-tuning/inference: length-sorted buckets + dynamic padding to a multiple of 8, boolean mask broadcastable to (B,1,1,L). Flash/mem-efficient backward is non-deterministic by default (PyTorch randomness notes via Context7) — set seeds and expect small run-to-run variance.
- **Version facts incidental to this work**: torch 2.14.1 is current on PyPI (requires_python >=3.10; cp314 manylinux_2_28 x86_64 wheel exists; depends on cuda-toolkit==13.0.3 wheels — driver 615.71 reports CUDA UMD 13.4, so compatible) — https://pypi.org/pypi/torch/2.14.1/json. torch 2.13.0 is the newest release in the ZeroGPU-supported list (cuda-toolkit 13.0.3, triton 3.7.1, cp310–cp314 wheels) — https://pypi.org/pypi/torch/2.13.0/json.
## 6. Context7 vs upstream
Context7 (/websites/pytorch_2_13) confirmed SDPBackend members (MATH, FLASH_ATTENTION, EFFICIENT_ATTENTION, CUDNN_ATTENTION), the 3-implementation note, `checkpoint(use_reentrant=False)` recommendation, and the autocast "promote"/unlisted-op rules, but did not return the autocast fp32 op list or SDPA constraints; those came from docs.pytorch.org 2.14 pages and the sdp_utils_cpp.h source as cited.
## 7. Calculator output (verbatim, `python3 model_math.py`)
```
Python 3.14.7
====================================================================================================
1) PARAMETER COUNTS (exact integers; M = 1e6)
====================================================================================================
sanity: BERT-base (post-LN, V=30522, +2 seg, +pooler) -> 109,482,240 (HF BertModel bert-base-uncased = 109,482,240)
config per-layer encoder embeddings MLM head(tied) TOTAL tied untied finetune N_nonemb
---------------------------------------------------------------------------------------------------------------------------------------------------------
A 24L d1024 h16 f4096 V32768 12,596,224 302,311,424 34,080,768 1,084,416 337,476,608 371.0M 337.4M 303.4M
B 24L d1024 h16 f4096 V32000 12,596,224 302,311,424 33,294,336 1,083,648 336,689,408 369.5M 336.7M 303.4M
C 24L d1024 h16 f4096 V16384 12,596,224 302,311,424 17,303,552 1,068,032 320,683,008 337.5M 320.7M 303.4M
D 20L d1024 h16 f4096 V32768 12,596,224 251,926,528 34,080,768 1,084,416 287,091,712 320.6M 287.1M 253.0M
E 22L d1024 h16 f4096 V32768 12,596,224 277,118,976 34,080,768 1,084,416 312,284,160 345.8M 312.3M 278.2M
F 24L d768 h12 f3072 V32768 7,087,872 170,110,464 25,560,576 624,896 196,295,936 221.5M 196.3M 170.7M
A' 23L d1024 h16 f4096 V32768 (extra option) 12,596,224 289,715,200 34,080,768 1,084,416 324,880,384 358.4M 324.8M 290.8M
config A per-layer breakdown: {'qkv': '3,148,800', 'out_proj': '1,049,600', '2xLN': '4,096', 'ffn_up': '4,198,400', 'ffn_down': '4,195,328', 'layer_total': '12,596,224'}
config A: embeddings = tok 33,554,432 + pos 524,288 + LN 2,048; MLM head tied = dense 1,049,600 + LN 2,048 + bias 32,768 = 1,084,416; classifier head = 1,051,650
distance from 330M (total, tied MLM head):
A' 23L d1024 h16 f4096 V32768 (extra option) 324.88M (-5.12M)
B 24L d1024 h16 f4096 V32000 336.69M (+6.69M)
A 24L d1024 h16 f4096 V32768 337.48M (+7.48M)
C 24L d1024 h16 f4096 V16384 320.68M (-9.32M)
E 22L d1024 h16 f4096 V32768 312.28M (-17.72M)
D 20L d1024 h16 f4096 V32768 287.09M (-42.91M)
F 24L d768 h12 f3072 V32768 196.30M (-133.70M)
vocab alignment: 32768 = 2^15 (mult of 64 and 128); 32000 = 64*500 = 128*250; 16384 = 2^14 -> all tensor-core aligned
====================================================================================================
2) ACTIVATION MEMORY MODEL (config A, L=512, bf16 autocast + fp32 master weights, SDPA)
====================================================================================================
residual stream kept in fp32 (what torch.autocast gives you: embeddings & residual adds stay fp32)
LN1 input = residual stream x 4,096 B = 4.00 d
QKV linear input (bf16 LN1 output) 2,048 B = 2.00 d
Q,K,V bf16 (saved by SDPA for backward) 6,144 B = 6.00 d
attn output O bf16 (SDPA bwd AND out-proj input; one tensor) 2,048 B = 2.00 d
SDPA logsumexp fp32 (one per head per query) 64 B = 0.06 d
out-proj dropout mask (bool, 1 B) 1,024 B = 1.00 d
LN2 input = residual stream x' 4,096 B = 4.00 d
FFN-up linear input (bf16 LN2 output) 2,048 B = 2.00 d
FFN-up output = GELU input bf16 (width f) 8,192 B = 8.00 d
GELU output = FFN-down input bf16 (width f) 8,192 B = 8.00 d
FFN-down dropout mask (bool, 1 B) 1,024 B = 1.00 d
TOTAL per token per layer 38,976 B = 38.06 d (c = 38 + 4h/d)
residual stream kept in bf16 (model.to(bf16) / Korthikanti fp16 assumption)
LN1 input = residual stream x 2,048 B = 2.00 d
QKV linear input (bf16 LN1 output) 2,048 B = 2.00 d
Q,K,V bf16 (saved by SDPA for backward) 6,144 B = 6.00 d
attn output O bf16 (SDPA bwd AND out-proj input; one tensor) 2,048 B = 2.00 d
SDPA logsumexp fp32 (one per head per query) 64 B = 0.06 d
out-proj dropout mask (bool, 1 B) 1,024 B = 1.00 d
LN2 input = residual stream x' 2,048 B = 2.00 d
FFN-up linear input (bf16 LN2 output) 2,048 B = 2.00 d
FFN-up output = GELU input bf16 (width f) 8,192 B = 8.00 d
GELU output = FFN-down input bf16 (width f) 8,192 B = 8.00 d
FFN-down dropout mask (bool, 1 B) 1,024 B = 1.00 d
TOTAL per token per layer 34,880 B = 34.06 d (c = 34 + 4h/d)
Korthikanti 2022 eager-attention reference: sbh(34 + 5*a*s/h) -> per token per layer = 34d + 5*a*s = 34,816 + 40,960 = 75,776 B -> the LxL score/softmax/dropout terms would be 118% ON TOP of 34d at L=512; SDPA flash/mem-efficient removes them (keeps only 4h B fp32 logsumexp).
under autocast an eager softmax runs in fp32, so the eager penalty would be >= 7*a*s, i.e. even worse than the fp16 paper figure
====================================================================================================
3) VRAM BUDGET (config A; GiB unless noted)
====================================================================================================
GPU: 23.99 GiB (25.76 GB) desktop already used: 1.90 CUDA ctx: 0.75 (0.5-1) headroom target: 4
static training state @16 B/param x 337,476,608 = 5.03 GiB (5.40 GB); + autocast bf16 weight cache 2 B/param = 0.63 GiB -> 5.66 GiB
=> left for activations+logits (incl. ~10% allocator slack): 11.68 GiB (training-process cap = 18.09 GiB)
per-sample (512 tok) layer activations: no-ckpt 457 MiB (c=38d) | 409 MiB (c=34d) ; full-ckpt fp32 inputs 48 MiB (+1 live layer 19.0 MiB) | bf16 inputs 24 MiB
per-sample logits+CE at peak (10 V B/pos): full 512 pos = 160 MiB ; masked ~15% (76 pos) = 24.0 MiB -> gather-before-head saves 136 MiB/sample (85%)
(i) NO checkpointing, c=38d (fp32 residual, autocast)
B logits | layers other logits acts+10% TOTAL headroom verdict
4 masked15% | 1.78 0.02 0.09 2.08 10.39 13.60 FITS
4 full512 | 1.78 0.04 0.62 2.63 10.94 13.05 FITS
8 masked15% | 3.57 0.04 0.19 4.16 12.47 11.52 FITS
8 full512 | 3.57 0.08 1.25 5.27 13.57 10.42 FITS
16 masked15% | 7.14 0.08 0.38 8.32 16.63 7.36 FITS
16 full512 | 7.14 0.16 2.50 10.53 18.84 5.15 FITS
24 masked15% | 10.71 0.13 0.56 12.48 20.78 3.20 tight
24 full512 | 10.71 0.25 3.75 15.80 24.10 -0.12 NO
32 masked15% | 14.27 0.17 0.75 16.64 24.94 -0.96 NO
32 full512 | 14.27 0.33 5.00 21.06 29.37 -5.38 NO
48 masked15% | 21.41 0.25 1.12 24.95 33.26 -9.27 NO
48 full512 | 21.41 0.49 7.50 31.59 39.90 -15.91 NO
(i') NO checkpointing, c=34d (bf16 residual)
B logits | layers other logits acts+10% TOTAL headroom verdict
4 masked15% | 1.60 0.02 0.09 1.87 10.18 13.81 FITS
4 full512 | 1.60 0.04 0.62 2.43 10.73 13.25 FITS
8 masked15% | 3.19 0.04 0.19 3.75 12.05 11.93 FITS
8 full512 | 3.19 0.08 1.25 4.85 13.16 10.83 FITS
16 masked15% | 6.39 0.08 0.38 7.49 15.80 8.19 FITS
16 full512 | 6.39 0.16 2.50 9.71 18.01 5.98 FITS
24 masked15% | 9.58 0.13 0.56 11.24 19.55 4.44 FITS
24 full512 | 9.58 0.25 3.75 14.56 22.87 1.12 NO
32 masked15% | 12.77 0.17 0.75 14.99 23.29 0.69 NO
32 full512 | 12.77 0.33 5.00 19.41 27.72 -3.73 NO
48 masked15% | 19.16 0.25 1.12 22.48 30.79 -6.80 NO
48 full512 | 19.16 0.49 7.50 29.12 37.42 -13.44 NO
(ii) FULL per-layer checkpointing, fp32 layer inputs (4d/tok/layer) + 1 live layer
B logits | layers other logits acts+10% TOTAL headroom verdict
4 masked15% | 0.26 0.02 0.09 0.40 8.71 15.28 FITS
4 full512 | 0.26 0.04 0.62 0.96 9.27 14.72 FITS
8 masked15% | 0.52 0.04 0.19 0.81 9.12 14.87 FITS
8 full512 | 0.52 0.08 1.25 1.92 10.22 13.76 FITS
16 masked15% | 1.05 0.08 0.38 1.62 9.93 14.06 FITS
16 full512 | 1.05 0.16 2.50 3.83 12.14 11.85 FITS
24 masked15% | 1.57 0.13 0.56 2.43 10.74 13.25 FITS
24 full512 | 1.57 0.25 3.75 5.75 14.06 9.93 FITS
32 masked15% | 2.09 0.17 0.75 3.24 11.55 12.44 FITS
32 full512 | 2.09 0.33 5.00 7.67 15.97 8.02 FITS
48 masked15% | 3.14 0.25 1.12 4.86 13.17 10.82 FITS
48 full512 | 3.14 0.49 7.50 11.50 19.81 4.18 FITS
(ii') FULL per-layer checkpointing, bf16 layer inputs (2d/tok/layer, task assumption)
B logits | layers other logits acts+10% TOTAL headroom verdict
4 masked15% | 0.16 0.02 0.09 0.29 8.60 15.39 FITS
4 full512 | 0.16 0.04 0.62 0.85 9.15 14.83 FITS
8 masked15% | 0.32 0.04 0.19 0.59 8.89 15.09 FITS
8 full512 | 0.32 0.08 1.25 1.69 10.00 13.99 FITS
16 masked15% | 0.64 0.08 0.38 1.17 9.48 14.51 FITS
16 full512 | 0.64 0.16 2.50 3.39 11.69 12.30 FITS
24 masked15% | 0.96 0.13 0.56 1.76 10.07 13.92 FITS
24 full512 | 0.96 0.25 3.75 5.08 13.39 10.60 FITS
32 masked15% | 1.28 0.17 0.75 2.35 10.65 13.33 FITS
32 full512 | 1.28 0.33 5.00 6.77 15.08 8.91 FITS
48 masked15% | 1.92 0.25 1.12 3.52 11.83 12.16 FITS
48 full512 | 1.92 0.49 7.50 10.16 18.46 5.52 FITS
====================================================================================================
4) TRAINING THROUGHPUT (config A, L=512 packed, B_eff=256 seq = 131,072 tok/step)
====================================================================================================
peak derivation: 71 TFLOPS (RTX 3090, BF16/FP32-acc dense) x 1.1241 (SM*clock ratio) = 79.8 TFLOPS; FP16-acc dense = 336 x 1.860 GHz x 256 = 160.0 TFLOPS -> half rate for FP32-acc = 80.0 TFLOPS; using 80
masked-only logits (15%): 6N=1.820 GF + attn 14*nl*d*L=0.176 GF + head=0.030 GF => 2.027 GFLOP/token (Kaplan/PaLM 12*nl*d*L = 0.151 GF; attention = 8.7% of total)
full-sequence logits: 6N=1.820 GF + attn 14*nl*d*L=0.176 GF + head=0.201 GF => 2.198 GFLOP/token (Kaplan/PaLM 12*nl*d*L = 0.151 GF; attention = 8.0% of total)
full checkpointing: 2.694 GFLOP/token = x1.329
peak = 80 TFLOPS dense bf16 (fp32-acc). MFU 30/37.5/45% -> 24/30/36 TFLOPS achieved
corpus steps/ep | tok/s @30% h/epoch | tok/s @37.5% h/epoch | tok/s @45% h/epoch | h/epoch ckpt@30% ckpt@45%
100M 763 | 11,843 2.35 | 14,804 1.88 | 17,764 1.56 | 3.12 2.08
150M 1,144 | 11,843 3.52 | 14,804 2.81 | 17,764 2.35 | 4.68 3.12
250M 1,907 | 11,843 5.86 | 14,804 4.69 | 17,764 3.91 | 7.79 5.20
400M 3,052 | 11,843 9.38 | 14,804 7.51 | 17,764 6.25 | 12.47 8.31
note: warmup of 1-2k steps exceeds one epoch at B_eff=256 for corpora <= 250M tokens -> plan multiple epochs or B_eff=128
====================================================================================================
5) INFERENCE LATENCY (fine-tuned classifier, batch 1, bf16 on GPU / fp32 on CPU)
====================================================================================================
L=512: 336.4 GFLOP -> 3090 Ti @30% MFU (24 TFLOPS): 14.0 ms | 5950X fp32 @1.0 TFLOPS eff: 336 ms | @0.5 TFLOPS: 673 ms (instructor benchmark 74 ms/sample)
L=256: 161.8 GFLOP -> 3090 Ti @30% MFU (24 TFLOPS): 6.7 ms | 5950X fp32 @1.0 TFLOPS eff: 162 ms | @0.5 TFLOPS: 324 ms (instructor benchmark 74 ms/sample)
L=128: 79.3 GFLOP -> 3090 Ti @30% MFU (24 TFLOPS): 3.3 ms | 5950X fp32 @1.0 TFLOPS eff: 79 ms | @0.5 TFLOPS: 159 ms (instructor benchmark 74 ms/sample)
5950X theoretical fp32 peak ~ 16 cores x 32 FLOP/cycle (2x256-bit FMA) x ~3.9 GHz = 2.0 TFLOPS; 1.0 TFLOPS effective = ~50%
GPU batch-1 is launch/latency-bound: 24 layers x ~12 kernels => several ms floor regardless of FLOPs (UNVERIFIED, measure)
```
## PINS
- torch==2.14.1 [local-training] — Current PyPI release (requires_python >=3.10) with a cp314-cp314-manylinux_2_28_x86_64 wheel for the system Python 3.14.7; depends on cuda-toolkit==13.0.3 runtime wheels, which the installed driver (CUDA UMD 13.4) satisfies. SDPA flash/mem-efficient backends and checkpoint(use_reentrant=False) are available. (https://pypi.org/pypi/torch/2.14.1/json)
- torch==2.13.0 [space] — Newest torch in the ZeroGPU supported list (2.8.0, 2.9.1, 2.10.0, 2.11.0, 2.12.1, 2.13.0); 2.14.x is not listed. PyPI confirms 2.13.0 exists with cp310-cp314 manylinux wheels (cuda-toolkit 13.0.3, triton 3.7.1). (https://huggingface.co/docs/hub/spaces-zerogpu)
- python==3.12 [space] — ZeroGPU supports only Python 3.12.12 and 3.10.13; the local 3.14.7 cannot be used for the Space, so keep model code/safetensors version-agnostic. (https://huggingface.co/docs/hub/spaces-zerogpu)
## RISKS
- The 80 TFLOPS bf16 peak is derived (whitepaper lists only the RTX 3090 at 71/142; the 3090 Ti figure is scaled by SM count x boost clock and cross-checked against Wikipedia's 320 sparse FP16-accumulate). Sustained clocks under the 450 W limit while sharing the GPU with the Wayland desktop will be lower, and eager PyTorch without torch.compile may land at 25-35% MFU rather than 45%; measure tokens/s on the first 50 steps and re-plan epochs.
- The activation model assumes the fused SDPA backends are actually used. Flash Attention rejects any non-null attn_mask (dispatcher warning verified in sdp_utils_cpp.h) and falls back to memory-efficient or, in the worst case, the math backend; a math fallback stores the LxL softmax (>= 7*a*s B/token/layer under autocast, more than doubling activation memory) and B=16 would no longer fit. Assert the backend with torch.nn.attention.sdpa_kernel and verify torch.cuda.max_memory_reserved() in the batch finder.
- c=38d assumes torch.autocast semantics (fp32 residual stream, fp32 LayerNorm inputs); if the implementation keeps a bf16 residual stream (c=34d) B=24 fits, if it accidentally keeps fp32 copies of linear inputs it will not. The 10% allocator slack is a guess; variable-length batches fragment more, so use PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True and pack/bucket to fixed shapes.
- Corpus size: 150M tokens is 1,144 optimizer steps per epoch at B_eff=256 (BERT used 1,000,000 steps); a 1-2k warmup exceeds one epoch for corpora <= 250M tokens. Expect an under-trained MLM unless several epochs are run or B_eff is reduced to 128; MLM quality limits are partly masked by the easy 2-class downstream task.
- From-scratch 24-layer training at peak LR >= 2e-4 can spike/diverge even with pre-LN; keep betas (0.9, 0.98), eps 1e-6, clip 1.0, and drop to 1e-4 if loss spikes. RoBERTa's 4e-4 was at batch 8K and does not transfer to batch 256.
- Sequence packing (RoBERTa FULL-SENTENCES) crosses email boundaries without cross-document masking; RoBERTa found DOC-SENTENCES slightly better. Acceptable for simplicity but note it in the report.
- CPU inference cannot meet the 74 ms/sample benchmark for a 337M model at 512 tokens (~336 ms at 1 TFLOPS effective, ~162 ms at 256 tokens); only a GPU (local 3090 Ti or ZeroGPU) or truncation to ~128 tokens gets there. ZeroGPU adds queue/cold-start latency, does not support torch.compile, and runs torch 2.8-2.13 on Python 3.10/3.12, so the deployment environment differs from the local 3.14 / torch 2.14.1 stack.
- Config A is 337.5M, 2.3% above the stated ~330M target; if the grader reads ~330M as a ceiling, switch to 23 layers (324.9M) before writing code, since depth changes every downstream number here by ~4%.
- Flash/memory-efficient SDPA backward passes are non-deterministic by default (PyTorch randomness notes); run-to-run metric variance should be expected and seeds documented.
- Unverified inputs carried from the brief: average email length of 200-300 tokens, the 5950X 1 TFLOPS effective fp32 rate, the ~1.9 GiB desktop usage (fluctuates with browser GPU use), and the batch-1 kernel-launch floor on the GPU; TechPowerUp could not be fetched (HTTP 403) and the HF bert-base parameter count used as a sanity anchor is from memory, although the formulas reproduce it exactly.
## SOURCES
- https://en.wikipedia.org/wiki/GeForce_RTX_30_series
- https://www.nvidia.com/content/PDF/nvidia-ampere-ga-102-gpu-architecture-whitepaper-v2.pdf
- https://ar5iv.labs.arxiv.org/html/2205.05198
- https://arxiv.org/pdf/2204.02311
- https://ar5iv.labs.arxiv.org/html/2001.08361
- https://ar5iv.labs.arxiv.org/html/2203.15556
- https://ar5iv.labs.arxiv.org/html/2205.14135
- https://ar5iv.labs.arxiv.org/html/1907.11692
- https://ar5iv.labs.arxiv.org/html/1810.04805
- https://ar5iv.labs.arxiv.org/html/2002.04745
- https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html
- https://huggingface.co/docs/hub/spaces-zerogpu
- https://docs.pytorch.org/docs/2.14/amp.html
- https://docs.pytorch.org/docs/2.14/notes/cuda.html
- https://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html
- https://docs.pytorch.org/docs/2.13/checkpoint.html
- https://docs.pytorch.org/docs/2.13/generated/torch.nn.attention.SDPBackend.html
- https://raw.githubusercontent.com/pytorch/pytorch/main/aten/src/ATen/native/transformers/sdp_utils_cpp.h
- https://github.com/Dao-AILab/flash-attention
- https://pypi.org/pypi/torch/json
- https://pypi.org/pypi/torch/2.14.1/json
- https://pypi.org/pypi/torch/2.13.0/json
- /tmp/claude-1000/-home-nathank-BITS-dev-Phishing-Model/8697d948-f234-4e8c-a1e6-c0b6f6ddd7ef/scratchpad/model_math.py
- /tmp/claude-1000/-home-nathank-BITS-dev-Phishing-Model/8697d948-f234-4e8c-a1e6-c0b6f6ddd7ef/scratchpad/model_math.out
- /tmp/claude-1000/-home-nathank-BITS-dev-Phishing-Model/8697d948-f234-4e8c-a1e6-c0b6f6ddd7ef/scratchpad/ga102.txt
- /tmp/claude-1000/-home-nathank-BITS-dev-Phishing-Model/8697d948-f234-4e8c-a1e6-c0b6f6ddd7ef/scratchpad/palm.txt
|