FEASIBILITY_AUDIT.md

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

Paper: *FlashOptim: Optimizers for Memory-Efficient Training* — arXiv [2602.23349v2](https://arxiv.org/abs/2602.23349), 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.