Cosmos3-Edge · single-GPU post-training

Fitting a 4B world-model fine-tune into 24 GB — without giving up training quality

NVIDIA's vision_sft_edge recipe for Cosmos3-Edge is written for 8×H100. Run as-is on one GPU it needs ~36 GiB. We traced every byte, proved which cuts are mathematically free, and landed the whole thing on a 22 GiB A10G — same effective batch, same EMA math, zero sample truncation.

Peak GPU memory

36 → 20 GiB

−44% vs stock recipe on one GPU

Effective batch parity

282 ≈ 283

samples/step vs stock 8-GPU config

EMA divergence vs stock

0.0

trainable-key A/B test, exact equality

VAE latents

bit-identical

offloaded vs resident, torch.equal

1 · The problem

Cosmos3-Edge is a 4.5B-parameter omni world model (3.36B transformer tower + Wan2.2 VAE + vision encoder). Its supervised fine-tuning recipe trains a 1.415B-parameter subset (keys_to_select: the generation-pathway MoE MLPs, time embedder, and modality bridges) on packed video sequences — designed and tested on 8×H100 80 GB. On a single 24 GB card (22.06 GiB usable) the stock configuration is ~14 GiB over budget, and the naive fixes — cutting sequence length, dropping EMA — are exactly the ones that threaten training quality.

The goal: make it fit without changing what the model learns. Every lever below is either provably equivalent to stock or was rejected.

2 · Where the memory actually goes

We read the training code path end-to-end (with independent adversarial re-verification of each claim) and reconciled it against live measurements. Hover any segment for details.

Component	Stock (est.)	L1 config (measured)	Full stack	What changed
Tower weights, bf16 (3.36B)	6.3	6.3	6.3	—
EMA (fp32)	12.5	0	0 (on CPU)	full-tower GPU clone → pinned-host subset buffers, same math
Adam m+v, bf16 (1.415B trainable)	5.3	5.3	5.3	—
Gradients, bf16 (trainable only)	2.6	2.6	2.6	—
Wan2.2 VAE, bf16	1.3	1.3	0 (on CPU)	moved to GPU only during encode; deterministic
CUDA context + buffers	2.4	2.4	2.4	—
Activations (full act-ckpt)	5.6 @45k tokens	3.4 @12k	3.4	packing cap 45,056 → 12,288; batch restored via grad-accum
Total	~36	21.3 measured	~20

3 · Three facts that make it free

These came out of the code audit; each was re-derived by an independent verification pass before we built on it.

4 · The design

Four levers, each zero-quality-risk. Both code changes are opt-in env flags — default off, stock behavior untouched.

Lever	Mechanism	GPU saved	Quality proof
L1 · packing cap + grad-accum	max_sequence_length 45,056 → 12,288; grad_accum_iter 2 → 61	~2–3 GiB	same effective batch; zero drops (packing simulation over the full dataset)
L2+L3 · CPU-resident subset EMA COSMOS_EMA_CPU_SUBSET=1	replaces the fp32 GPU net_ema clone with fp32 pinned-host buffers over trainable params; same foreach lerp, same beta schedule; checkpoint save reconstructs the full net_ema.* dict exactly	12.5 GiB	10-step A/B vs stock updater: trainable-key diff 0.0; save→load round-trip identical
L4 · VAE CPU round-trip COSMOS_VAE_CPU_OFFLOAD=1	frozen Wan2.2 VAE parked on CPU; moved to GPU only inside encode/decode	1.3 GiB	encode is deterministic (mean-only, no sampling): latents bit-identical, +0.4 s/call

Rejected on evidence (the anti-levers)

• FusedAdam / Muon / Dion2 — all force fp32 optimizer state (+5–10 GiB, not −).
• FSDP2 CPU-offload — documented in-repo as never wired to any recipe; single-GPU FSDP doesn't shard anyway.
• Optimizer-in-backward — silently disables the recipe's global grad clip (0.1) and conflicts with grad accumulation.
• mxfp8/nvfp4 quantization config — inference-only PTQ; no training path.
• Cutting max_caption_tokens — the one knob that does truncate (max caption = 1,841 tokens; the 2,048 cap is lossless — leave it).

5 · Implementation

Everything lives on the feat/24gb-single-gpu-sft branch: two guarded code paths (+679/−22 lines), three recipe TOMLs, and launchers.

The heart of it — the CPU subset EMA update (same lerp the stock GPU worker runs):

# cosmos_framework/utils/generator/cpu_subset_ema.py
@torch.no_grad()
def update_average(self, net, beta):
    """ema = beta * ema + (1 - beta) * param  — identical math to the stock worker."""
    names = []
    for name, p in net.named_parameters():
        if name in self._buffers:                     # trainable subset only
            self._staging[name].copy_(dt2lt(p).detach(), non_blocking=True)  # async D2H
            names.append(name)
    torch.cuda.synchronize()
    targets = [self._buffers[n] for n in names]       # fp32, pinned host memory
    sources = [self._staging[n].to(torch.float32) for n in names]
    torch._foreach_mul_(targets, beta)
    torch._foreach_add_(targets, sources, alpha=1.0 - beta)

At checkpoint time the full net_ema.* state dict is reconstructed exactly: EMA buffers for trainable keys, fp32 casts of the live net for frozen keys (equal by definition). Save→load round-trips are verified identical, and ema_scope (used by validation/sampling) swaps the subset in and out losslessly.

Run it

git clone -b feat/24gb-single-gpu-sft https://github.com/linjiw/cosmos-framework
# ...standard cosmos-framework setup (uv sync --group cu128), stage dataset + DCP ckpt...

COSMOS_EMA_CPU_SUBSET=1 COSMOS_VAE_CPU_OFFLOAD=1 \
NPROC_PER_NODE=1 WAN_VAE_PATH=$WAN \
  bash examples/launch_sft_vision_edge_24gb.sh
# variants: MEMTEST=1 (fast memory smoke) · EMAVAL=1 (EMA-correctness config)

6 · Validation status

Gate	Method	Result
Config run trains sanely	18 optimizer steps @ cap 12,288 on an empty GPU	PASS — loss 2.1–3.6, finite grads, 0 drops
Peak memory measured	2 s VRAM sampler across the full run	21.3 GiB (L1 only) — above the 19.5 estimate; backward transient at 12k is real
VAE offload equivalence	encode a fixed clip, resident vs offloaded	PASS — torch.equal true; +0.38 s/round-trip
EMA equivalence (unit)	10-step A/B vs stock DTensorFastEmaModelUpdater	PASS — trainable keys exactly equal; frozen keys exact (stock itself drifts 7e-9 lerping constants)
EMA save/load round-trip	state-dict reconstruction + reload	PASS — identical buffers
Integrated run, EMA on + both flags	smoke train with all levers active	queued — waits for a co-tenant-free window on our shared box
500-iter run + sample-quality eval	full fine-tune, export, i2v eval vs base	pending the above