wfen/Cosmos3-Nano-NVFP4

Cosmos3-Nano — NVFP4-AWQ (safetensors)

Weight-only NVFP4 (E2M1, block 16) + AWQ quantization of the Cosmos3OmniTransformer for Cosmos3-Nano, delivered as safetensors. The transformer drops from 30 GB (bf16) to **9 GB**; VAE, vision encoder, tokenizers, and scheduler remain bf16. Runs the diffusers Cosmos3OmniPipeline (fake-quant) on a single RTX 5090 (32 GB, sm_120).

Load

from load_quantized import load           # self-contained; needs torch, diffusers, modelopt, safetensors
pipe = load(".")                           # this directory; requires a Blackwell GPU at load (FP4 capability check)
import torch
with torch.autocast("cuda", torch.bfloat16):
    img = pipe("a corgi astronaut", num_frames=1, height=480, width=480).video[0][0]
img.save("out.png")

Format (Path B)

The transformer is serialized as safetensors plus a tiny structural sidecar:

File	Contents
transformer/diffusion_pytorch_model.safetensors	505 weight-only NVFP4 weights + scales + the AWQ smoothing scales + bf16 keep-modules (≈2829 tensors)
transformer/modelopt_state.pt	tensor-free ModelOpt structural state (quantizer layout) — needed to rebuild the quantizer modules
transformer/config.json	transformer config (action_gen=false)
quantization_config.json	recipe, exclusions, and a scale_layout summary (see the layout note below)

Load = from_config (action_gen=False) → modelopt.torch.opt.restore_from_modelopt_state → load_state_dict(strict=True), on CPU, then move the whole pipeline to the GPU. The loader reads only the safetensors + sidecar — never the .pt. NVFP4 restore is device-order-sensitive (CPU-restore-then-.to(device)) and requires a Blackwell GPU at load time (FP4 capability check).

NVFP4 scale layout (authoritative; per quantized Linear)

tensor	shape (example, in=4096)	dtype	role
…weight	(out, in/2)	uint8	packed FP4 E2M1 (2 codes / byte)
…weight_quantizer._scale	(out, in/16)	E4M3	per-block-16 scale
…weight_quantizer._double_scale	()	float32	per-tensor FP32 global scale
…weight_quantizer._amax	()	bf16	per-tensor amax
…input_quantizer._pre_quant_scale	(in,)	bf16	AWQ smoothing (weight-only)

Recipe & scope (INV-2 / INV-4)

Weight-only NVFP4 (E2M1, block 16, per-block E4M3 scale + per-tensor FP32 global) with AWQ calibration (awq_lite; activation quantizers disabled). Quantized: self_attn.*, mlp.*, mlp_moe_gen.*, lm_head (505 Linears). Kept bf16: token embeddings, all norms, time_embedder, proj_in/proj_out, audio/action adapters. Calibration: 64 prompts (3 bundled + 61 curated physical-AI), 8 denoising steps, seed 123.

Quality

Verified by equivalence to the reference NVFP4 checkpoint on EC-01..06 against the frozen Session-1 band (M2 LPIPS ≤ τ_nvfp4 = 0.5463 and mean SSIM ≥ 0.90, no catastrophic artifacts), plus a bitwise safetensors==.pt weight round-trip. NVFP4 ≠ reference bitwise by design (our AWQ calibration set differs from the reference's unshipped set).