FEASIBILITY_AUDIT.md

Feasibility audit & claim-to-evidence plan — FlashOptim (ICML 2026 Spotlight)

Paper: FlashOptim: Optimizers for Memory-Efficient Training — arXiv 2602.23349v2, OpenReview Wfe1iJocjF Authors: Gonzalez Ortiz, Gupta, Rinard, Blalock (Databricks AI Research) Code: https://github.com/databricks/flashoptim (Apache-2.0, v0.1.4, ~3.2k LoC single module + 6.1k LoC tests) Date of audit: 2026-07-15. Reproducer: ProCreations (HF), Claude Code session.

1. What the paper claims (assigned challenge claims)

1. Mechanism (§3): FlashOptim = improved master-weight splitting (ULP-normalized error correction, Alg. 1) + companded 8-bit optimizer-state quantization (softsign for momentum, sqrt for variance; Alg. 2–3), integrated in a prologue/epilogue around standard updates (Alg. 4–6).
2. Bytes/param (Table 1): AdamW training memory 16 → 7 bytes/param (4+4+4+4 → 2 θ′ + 1 ρ + 2 grad + 1 m + 1 v), or 5 with gradient release (grads not materialized). Scale-factor overhead (2/32 bytes ×2) acknowledged in §4.3 but omitted from Table 1.
3. Llama-3.1-8B finetune (Fig. 1, Table 4): peak 175.2 → 112.9 GiB (−36%); Params 29.9 → 15.0 GiB (−50%); Optim 59.8 → 23.4 GiB (−61%); step 12.5 → 11.5 ms. Setup: FSDP2 + activation ckpt, H100s, 5.2M tokens/step.
4. Convergence (Fig. 2): FlashAdamW matches AdamW loss trajectory on GPT-2 124M pretraining (20k steps, FineWeb10B); FlashSGD matches SGD on ResNet-50/ImageNet (90 epochs).
5. Final scores (Tables 2–3): ImageNet top-1 (SGD 77.01±.02 vs 77.16±.09; AdamW 75.51±.09 vs 75.67±.04), GSM8K after 8B finetune (75.09±.40 vs 74.98±.77), GPT-2 ICL suite (8 benchmarks, means 44.0±.4 vs 45.0±1.0 AdamW; 44.0±.5 vs 44.7±.5 Lion). 3 seeds each.
6. Ablations (Fig. 3, Fig. 5): ULP-based splitting gives much lower FP32 reconstruction error than BF16+BF16 error storage (BF16+INT16 <1e-9 mean rel. err., 99.92% bitwise-exact; BF16+BF16 >1e-6). Companding is necessary: linear 8-bit quantization of AdamW states diverges in GPT-2 training; companded is stable.

2. Artifact inventory

Artifact	Status	Notes
Paper PDF/HTML v2	✅ obtained	full method pseudocode (Alg. 1–6), all hyperparams in App. B
Code (GitHub)	✅ cloned	optimizers + tests only. No paper experiment scripts (no ImageNet/nanoGPT/Llama recipes, no Fig. 3/4/5 scripts)
Reference optimizers	✅ in test/reference.py	plain PyTorch, device-agnostic
Test suite	✅ 6.1k LoC	GitHub CI runs lint only — tests are NOT publicly CI-verified; running them is real audit value
Checkpoints/data	❌ none released	ImageNet gated ~150GB; FineWeb10B ~20GB; both exceed download ceiling
OpenReview reviews	⚠️ API bot-blocked	not essential

Hard constraint discovered: flashoptim imports triton unconditionally and raises NotImplementedError for quantization / ECC / fused steps without CUDA (optimizers.py:485-492). → The authors' library runs only on Linux + NVIDIA GPU. Local macOS work must use independent reimplementations of Alg. 1–3 (pure NumPy/PyTorch); the authors' code runs in HF GPU Jobs.

Key API facts (verified in source): master_weight_bits ∈ {24, 32, None}; quantize: bool; fused: bool; per-state QuantizedTensorSpec(signed, sqrt, softsign) — FlashAdam uses exp_avg: signed+softsign, exp_avg_sq: unsigned+sqrt (exactly Alg. 2–3). Subclassing _quantized_state_spec to disable sqrt/softsign yields the paper's "linear quantization" baseline for the Fig. 5 divergence test. enable_gradient_release(model, opt) exists. compute_ecc_bits/reconstruct_fp32_param expose Alg. 1.

Issue tracker findings: Issue #1 (closed): naive benchmarking (3 variants in one process) shows FlashOptim worse; author prescribes one-variant-per-process methodology — we adopt it. Issue #2 (closed): master_weight_bits=None + bf16 params = pure-bf16 training, known-unstable. Issue #6 (open): FSDP2 uneven-shard checkpoint bug — irrelevant to our single-GPU scope.

3. Local environment

• macOS (Apple Silicon), no CUDA → no authors' code locally. 16 GiB free disk. PyTorch available locally (checked at smoke-test time).
• HF: authenticated (ProCreations), 0 active jobs, hard ceiling $2.00 this paper.
• GPU flavor choice: l4x1 $0.80/h (Ada sm_89: full bf16 + Triton ≥3 support, 24 GB). T4 rejected (Turing: no reliable bf16 paths). a10g-small ($1.00/h) is fallback if L4 queues/fails.

4. Claim-to-evidence matrix

Evidence classes: DR direct reproduction · INC independent numerical check · RAA released-artifact audit · SE scaled experiment (GPU, reduced scale, mechanism intact) · TOY toy proxy · UN unavailable.

#	Claim	Evidence plan	Class	Where	Est. cost
1	Mechanism §3	(a) Code audit mapping Alg. 1–4 to source lines; (b) independent CPU reimplementation of Alg. 1–3; (c) equivalence test: my eager impl vs authors' Triton kernels, elementwise on GPU; (d) authors' test suite run	RAA + INC	local CPU + Job A	$0 + shared Job A
2	16→7 (5⋆) bytes/param	(a) Analytic accounting from dtypes (deterministic); (b) measured torch.cuda.memory_allocated deltas per tensor class across model sizes 5M–400M, slope in bytes/param, AdamW vs FlashAdamW vs +gradient-release, one variant per process	INC + SE	Job B	shared Job B
3	8B: 175→113 GiB peak	(a) INC: recompute Table 4 Params/Optim rows exactly from 8.03e9 params × dtype widths; (b) SE: extrapolate measured bytes/param slope (from claim 2) to 8.03B and compare vs Table 4; (c) full 8B run UN under $2 (needs ≥8×H100) — scope labeled honestly. Peak-GiB number depends on activations/batch → only consistency-checkable	INC + SE + UN(full)	local + Job B	$0
4	Loss trajectories match (Fig. 2)	SE: small-GPT (~10–20M) on real text, AdamW vs FlashAdamW, identical seed/data order, overlaid loss curves + final-loss delta; CNN on CIFAR-10/MNIST, SGD vs FlashSGD. Reduced scale labeled (paper: 124M/20k steps; ResNet-50/ImageNet)	SE	Job B	shared
5	Final scores match (Tables 2–3)	Full ImageNet-90ep / 8B-GSM8K / 10B-token-ICL: UN under budget (≫$100). TOY: small-model final val-acc parity (≥2 seeds) CIFAR-10; plus RAA-style statistical audit of reported Tables 2–3 (σ-overlap analysis). Labeled toy, not presented as full repro	TOY + UN(full)	Job B + local	shared
6	Fig. 3 recon error; Fig. 5 divergence	(a) INC/DR: exhaustive FP32 reconstruction-error sweep (all 2^32 bitstrings, same protocol as paper) comparing baseline/BF16+BF16/ULP+INT8/ULP+INT16, per-exponent mean rel. error + bitwise-exact %; (b) NMSE companded-vs-linear on real harvested optimizer states (Fig. 4-style); (c) SE: LM training w/ authors' code, companded vs linear (spec override) → divergence check	INC/DR + SE	local CPU + Job B	$0 + shared

5. Costed GPU plan — REVISED 2026-07-16 per updated orchestration rules

Budget revision: paper ceiling raised $2.00 → $4.00; hardware restricted to one T4 or A10G-small (L4 prohibited); max 3 paid Jobs, one active at a time. Quality steering (QUALITY_STEERING.md) requires: authors' Triton kernels + tests vs the independent eager implementation; one-variant-per-process memory slopes across sizes ± gradient release; real-text GPT at largest defensible scale (target GPT-2 124M) with matched arms; companded-vs-linear divergence diagnostics; ≥2 seeds or paired trajectory analysis.

Job	Content	Timeout	Upper-bound cost (a10g-small $1.00/h)
1 (audit+smoke)	pip install authors' repo; full non-distributed test suite; eager-vs-Triton equivalence (Alg. 1–3); tiny-scale dry-run of Job 2 scripts; record kernel/driver versions	30 min	$0.50
2 (main)	(a) memory-scaling sweep, one variant per process, sizes ~6M→350M × {ref AdamW, FlashAdamW, +grad-release, split-only, quant-only}; (b) GPT-2 124M on FineWeb10B shards (kjj0/fineweb10B-gpt2, ~200 MB/shard ≤ download cap): AdamW s0 vs FlashAdamW s0 (identical data/init), AdamW s1 seed-noise yardstick, linear-quant divergence arm	≤170 min	$2.84
3 (contingency)	only for a proven infrastructure/protocol defect per rules	≤40 min	$0.67
Total			planned ≤$3.34, worst-case $4.00 = ceiling

Rules honored: no H100/multi-GPU/L4; bounded --timeout; labels paper=Wfe1iJocjF; status checks after submit; cumulative ledger in logbook.

6. Deliberate scope exclusions (honesty ledger)

• No Llama-3.1-8B finetune, no ImageNet-1K training, no 10B-token pretraining, no GSM8K eval of a finetuned 8B — each alone exceeds the total budget by orders of magnitude.
• ICL benchmark suite (Table 3) not run: models we can train under $2 (≤20M params, ≤100M tokens) score at chance on these benchmarks, so parity would be vacuous — worse than no evidence.
• Multi-GPU (FSDP2/DDP) paths untested; single-GPU only.
• Step-time (ms) comparisons reported only qualitatively: L4 ≠ H100.