wfen/Cosmos3-Nano-FP8-Blockwise

Cosmos3-Nano FP8 Blockwise Mixed-Precision Checkpoint

Mixed-precision blockwise FP8 weight-only quantization of Cosmos3OmniTransformer.

Recipe

• Quantized (blockwise FP8 E4M3, 128x128): mlp, mlp_moe_gen, lm_head (217 modules)
• Kept bf16: self_attn, embed_tokens, all norms, time_embedder, proj_in/proj_out, audio/action adapters
• Activations: bf16 (weight-only quantization)
• Algorithm: max (no calibration data needed)
• Framework: NVIDIA ModelOpt 0.44.0 + safetensors export

Quick Start

# Verify the checkpoint loads correctly
python load_checkpoint.py --verify

# Generate a single frame (smoke test)
python load_checkpoint.py \
  --prompt "A robotic arm in a kitchen" \
  --steps 8 --frames 1

# Generate a multi-frame video (quality)
python load_checkpoint.py \
  --prompt "A robotic arm in a kitchen" \
  --steps 35 --frames 57 \
  --height 480 --width 640

Dependencies

• Python 3.12+
• PyTorch 2.11+ with CUDA support
• diffusers 0.39.0+
• modelopt 0.44.0+
• safetensors

All available in the project's Docker environment.

Checkpoint Contents

transformer/
  config.json                              # Cosmos3OmniTransformer config (action_gen=False)
  diffusion_pytorch_model.safetensors      # FP8 weights + blockwise scales (~18.8 GB)
  modelopt_state.pt                        # Structural sidecar (~670 KB, quantizer topology)
quantization_config.json                   # Recipe, block size, scale layout documentation
quantizer_map_diff.json                    # INV-2 validation result
load_checkpoint.py                         # Standalone loader
README.md                                  # This file

The modelopt_state.pt sidecar contains only the quantizer structure (which modules have quantizers, their configs). It does NOT contain model weights. It uses pickle format (weights_only=False on load) and should only be trusted from this locally-produced checkpoint.

Loading Programmatically

import torch, glob
import modelopt.torch.opt as mto
from diffusers import Cosmos3OmniPipeline, Cosmos3OmniTransformer, UniPCMultistepScheduler
from safetensors.torch import load_file

CKPT = "dist/Cosmos3-Nano-FP8-Blockwise"

# 1. Build skeleton
cfg = {**Cosmos3OmniTransformer.load_config(f"{CKPT}/transformer/config.json"), "action_gen": False}
transformer = Cosmos3OmniTransformer.from_config(cfg).to(torch.bfloat16)

# 2. Restore quantizer structure from sidecar
state = torch.load(f"{CKPT}/transformer/modelopt_state.pt", weights_only=False)
restored = mto.restore_from_modelopt_state(transformer, state)
if restored:
    transformer = restored

# 3. Load weights + scales from safetensors
tensors = {}
for shard in sorted(glob.glob(f"{CKPT}/transformer/*.safetensors")):
    tensors.update(load_file(shard))
transformer.load_state_dict(tensors, strict=True)

# 4. Build pipeline
pipe = Cosmos3OmniPipeline.from_pretrained(
    CKPT, transformer=transformer, torch_dtype=torch.bfloat16, enable_safety_checker=False
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=10.0)
pipe = pipe.to("cuda")

# 5. Generate under autocast
with torch.autocast("cuda", torch.bfloat16):
    result = pipe(prompt="...", num_frames=57, height=480, width=640, num_inference_steps=35,
                  generator=torch.Generator("cpu").manual_seed(123))

Quality Summary

Compared to Phase 1 per-tensor FP8 (vs bf16 gold standard):

Case	Improved?	LPIPS Delta
EC-01 (t2v)	No	-0.056
EC-02 (sound/MoE)	Yes	+0.069
EC-03 (i2v)	Yes	+0.010
EC-05 (hard)	Yes	+0.029
EC-06 (OOD)	Yes	+0.025

4/6 cases improved (DoD-3 PASS). See docs/reports/phase_2_quality_report.md for the full comparison.

Generation Parameters (INV-5 Determinism)

For reproducible output, use these exact settings:

• Seed: 123
• Scheduler: UniPCMultistepScheduler(flow_shift=10.0)
• CUBLAS: CUBLAS_WORKSPACE_CONFIG=:4096:8
• Autocast: torch.autocast("cuda", torch.bfloat16)
• Generator device: "cpu" (not "cuda")
• Steps: 35 (quality) or 8 (smoke)

Generated Examples

In folder 'assets/FP8-Examples', you can find a selection of generated videos from the FP8 blockwise checkpoint. Each example includes VRAM reports and ffprobe reports.