Hugging Face's logo

BuildITSystems
/
phishing-detector-encoder

Text Classification
PyTorch
Safetensors
English
phishing
email-security
from-scratch
spam-detection
Eval Results (legacy)
Model card Files
xet
 Community
phishing-detector-encoder / presentation /docs /research /research_model-math.md
BuildITSystems's picture
BuildITSystems
Add presentation bundle (reports, curves, run log, code) for the Colab notebook
439e7b1 verified
preview code
|
Raw History Blame Contribute Delete
39.4 kB

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 (>= 7as 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