---
license: openmdw-1.0
base_model:
- nvidia/Cosmos3-Nano
---
# 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

```bash
# 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

```python
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.