Cosmos3-Edge under 10 GB

Quantizing and optimizing NVIDIA's 4B omni world model to run every inference mode in under 10 GiB of GPU memory on an Ampere A10G — with quality validated against a properly calibrated gate, and a torch 2.10 kernel trap found on the way.

July 2026 · A10G 23 GB (sm_86) · torch 2.10 + torchao 0.16 · repo · framework patch · previous phase: SFT in 24 GB

• 1 · Why this is not just "run bitsandbytes"
• 2 · Where the bytes actually are
• 3 · What we changed
• 4 · Memory results
• 5 · Evaluation: quality without fooling ourselves
• 6 · The torch 2.10 kernel trap (12× slowdown, root-caused)
• 7 · Speed results
• 8 · The recipe
• 9 · What's next

1 · Why this is not just "run bitsandbytes"

Cosmos3-Edge is not a plain LLM. It is a Mixture-of-Transformers: one 28-layer tower in which every layer carries two disjoint parameter sets — an understanding/reasoner pathway (Nemotron-2B VL) and a generation pathway (a rectified-flow video diffusion expert) — sharing a joint attention. Both pathways run on every diffusion step; there is no way to load just one. Add the Wan2.2 video VAE and a SigLIP2 vision tower and you get ~8.5 GiB of bf16 weights before a single activation.

NVIDIA ships quantization hooks in cosmos-framework — but the two shipped methods (mxfp8, nvfp4) require Blackwell tensor cores. On the Ampere A10G (sm_86: no FP8, no FP4) they cannot run at all. NVIDIA's own Jetson benchmarks for Edge run plain bf16 on 16–128 GB unified-memory devices. As far as we could find, no sub-10 GB Cosmos3-Edge recipe existed anywhere — this write-up is one.

We also found the shipped hooks were never applied on the checkpoint path most people actually use (the HuggingFace/diffusers snapshot): the config was stored into the model dict and silently ignored. Part of the patch is simply making --quantization-method work there.

2 · Where the bytes actually are

We parsed the safetensors headers for exact counts (no estimates):

Two structural facts drive everything. First, at 480p×121-frame video the diffusion runs in a 16×/4× compressed latent space — only ≈12k tokens per forward — so weights, not activations, dominate the peak: quantize weights and the peak follows. Second, classifier-free guidance runs as two sequential full forwards per step (70 forwards per 35-step video), which is why the eventual distillation track (4-step, CFG-free) is worth ~17× compute.

3 · What we changed

Around the patch we built a measurement harness (bench/): phase-separated memory probes (load transient vs inference peak — a distinction that turned out to be gate-critical), a fixed 4-mode battery with every sampler argument pinned, PSNR/LPIPS comparison, and a sampling-trajectory probe for the quality gate below. The design was adversarially reviewed by two independent review passes (one for code-path correctness — every claim verified against file:line — one for experimental methodology) before we trusted the numbers.

4 · Memory results

The −2.62 GiB delta is identical in every mode and matches the arithmetic (2.82 B params × 1 byte saved, minus scales) to ±0.02 — evidence that int8wo adds no activation overhead: the whole saving survives to the peak. Stacking the existing COSMOS_VAE_CPU_OFFLOAD=1 flag frees another 1.31 GiB at the 480p peak (≈8.6 GiB expected steady-state).

Honest caveat: a hard-capped proof run (allocator clamped to ~10.1 GiB via DEVICE_MEMORY_FRACTION) survived load but OOMed on a decode transient at 9.53 GiB — the last ~0.5 GiB of transient headroom at 480p is open work (§9). The un-capped 480p peak of 9.86 GiB stands.

5 · Evaluation: quality without fooling ourselves

The obvious evaluation — same seed, compare pixels — is wrong for diffusion models, and it's worth being precise about why. A 480p video is 50 sampling steps × 2 CFG forwards = 100 forwards; a 0.05% numeric difference per forward compounds into a different sample — different but equally good. Per-pixel PSNR measures trajectory divergence, not quality loss. Our gate had to separate those.

Tier 0 — the null floor

Two identical bf16 runs, same seed: bit-exact (relative L2 = 0.0). The pipeline is deterministic, so every non-zero difference we measure is real signal, not run-to-run noise.

Tier 1 — direct quantization error, before it compounds

We instrumented the UniPC sampler and captured the step-0 velocity prediction — at step 0 the inputs (noise, timestep, prompt) are bit-identical across runs, so the difference there is the pure single-forward quantization error. The calibration insight: merely toggling torch.compile on the bf16 model — same weights, universally accepted as "the same model" — shifts the step-0 velocity by 10.1% relative L2. That's the numerics delta everyone already accepts. Quantization must be judged against it.

Tier 2 — perceptual, on full generations

The reasoner mode decodes greedily; int8 flips occasional near-tie tokens, after which the texts legitimately diverge — both runs produce valid, complete task plans for the same manipulation prompt. (The framework's own docstrings document the same greedy-tie divergence for bf16 rank-to-rank noise.) A teacher-forced token-agreement probe is the planned tightening here.

6 · The torch 2.10 kernel trap

First int8 run: memory perfect, sampling 12× slower. No graph breaks, nothing in the logs. The root cause is worth publishing on its own:

7 · Speed results

Weight-only int8 being fast is not luck: it halves weight-memory traffic, and these forwards are partly bandwidth-bound even at 12k tokens. The clearest demonstration is the reasoner mode: 51.2 s → 42.4 s (17% faster) under int8 — autoregressive decode is memory-bound, exactly where weight-only quantization pays twice (memory and speed).

8 · The recipe

# apply the patch branch to cosmos-framework, then:
python -m cosmos_framework.scripts.inference \
    -i input.json -o outputs/run \
    --checkpoint-path Cosmos3-Edge \
    --quantization-method int8wo \
    --no-use-torch-compile          # eager int8 ≈ compiled-bf16 speed, no inductor footguns

# stack for the lowest peak (−1.31 GiB more):
COSMOS_VAE_CPU_OFFLOAD=1 python -m cosmos_framework.scripts.inference ...

Every number on this page is reproducible from the repo: bench/run_battery.sh bf16_ref → bench/run_battery.sh int8wo --quantization-method int8wo → bench/compare.py bf16_ref int8wo; trajectory gates via bench/trajprobe.py; the kernel regression isolated in bench/int8wo_microbench.py. Raw measurement JSON is in results/.

9 · What's next

The full research document — literature review across quantization/attention/offload/ distillation (60+ sources), the verified model anatomy, ranked technique stacks with memory arithmetic, and the review-hardened measurement protocol — is in docs/QUANT_10GB_RESEARCH_PLAN.md.