Skip to content

cookbook(diffusers): add FP8 quantized checkpoint support #304

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ConstBob wants to merge 3 commits into
base: main
Choose a base branch
from
Open

cookbook(diffusers): add FP8 quantized checkpoint support #304

Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
39b17e0
cookbook(diffusers): add FP8 quantized checkpoint loading path
ConstBob Aug 5, 2026
8d34c6c
Merge branch 'NVIDIA:main' into feat/diffusers-fp8
ConstBob Aug 6, 2026
f4878b6
cookbook(diffusers): fix Distilled FP8 ModelOpt restore via transform…
ConstBob Aug 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
115 changes: 104 additions & 11 deletions cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
"\n",
"Run the Cosmos3-Nano, Cosmos3-Super, or Cosmos3-Edge examples independently. Cosmos3-Edge has no audio modules and uses its 480p generation settings. Run Cosmos3-Super T2V/I2V examples without audio. Each section loads the matching model explicitly.\n",
"\n",
"**FP8:** `Cosmos3-Nano`, `Cosmos3-Super`, and the Super 4-Step distilled checkpoints also ship static-scale FP8 builds. Enable them with `COSMOS3_DIFFUSERS_PRECISION=fp8` (see **Quantized Checkpoints** below). Request cells stay the same; only checkpoint loading changes. `Cosmos3-Edge` has no FP8 build.\n",
"\n",
"Note: if you have already completed steps 1-3 and installed the `Cosmos3 Diffusers (Python 3.13)` kernel, switch to that kernel and jump directly to step 4. Run the restore cell there, then continue with verification and the examples.\n"
]
},
Expand Down Expand Up @@ -57,6 +59,8 @@
"export HF_HOME=/path/to/large/huggingface/cache\n",
"export UV_LINK_MODE=copy\n",
"export CUDA_VISIBLE_DEVICES=0\n",
"# Optional: bf16 (default) or fp8 for Nano/Super/4-Step — see Quantized Checkpoints\n",
"export COSMOS3_DIFFUSERS_PRECISION=bf16\n",
"```\n"
]
},
Expand Down Expand Up @@ -103,6 +107,7 @@
" os.environ.setdefault(\"HF_HOME\", str(Path.home() / \".cache\" / \"huggingface\"))\n",
" os.environ.setdefault(\"HF_HUB_DISABLE_XET\", \"1\")\n",
" os.environ.setdefault(\"CUDA_VISIBLE_DEVICES\", \"0\")\n",
" os.environ.setdefault(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\")\n",
"\n",
" print(f\"COSMOS_ROOT: {COSMOS_ROOT}\")\n",
" for key in [\n",
Expand All @@ -114,6 +119,7 @@
" \"HF_HOME\",\n",
" \"HF_HUB_DISABLE_XET\",\n",
" \"CUDA_VISIBLE_DEVICES\",\n",
" \"COSMOS3_DIFFUSERS_PRECISION\",\n",
" ]:\n",
" print(f\"{key}: {os.environ[key]}\")\n",
" print(\"HF_TOKEN:\", \"<set>\" if os.environ.get(\"HF_TOKEN\") else \"<unset>\")\n",
Expand Down Expand Up @@ -159,7 +165,8 @@
" ipykernel \\\n",
" torch \\\n",
" torchvision \\\n",
" transformers\n",
" transformers \\\n",
" \"nvidia-modelopt==0.44.0\"\n",
"\n",
"\"$COSMOS3_DIFFUSERS_VENV/bin/python\" -m ipykernel install --user \\\n",
" --name cosmos3-diffusers \\\n",
Expand Down Expand Up @@ -631,10 +638,38 @@
" \"Cosmos3-Super-Image2Video-4Step\",\n",
"}\n",
"\n",
"# FP8 builds for Nano/Super and Super 4-Step (same request path as bf16; see Quantized Checkpoints).\n",
"FP8_SUPPORTED_MODELS = {\n",
" \"Cosmos3-Nano\",\n",
" \"Cosmos3-Super\",\n",
" \"Cosmos3-Super-Text2Image-4Step\",\n",
" \"Cosmos3-Super-Image2Video-4Step\",\n",
"}\n",
"\n",
"_pipe = None\n",
"_pipe_model = None\n",
"\n",
"\n",
"def precision_mode() -> str:\n",
" return os.environ.get(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\").strip().lower()\n",
"\n",
"\n",
"def is_fp8_precision() -> bool:\n",
" return precision_mode() == \"fp8\"\n",
"\n",
"\n",
"def resolve_fp8_pretrained_source(model: str):\n",
" \"\"\"Return a local path or public HF repo id for the FP8 checkpoint.\n",
"\n",
" Uses `MODEL_IDS[model]` with `revision=COSMOS3_FP8_REVISION` (default `fp8`) unless\n",
" `COSMOS3_FP8_MODEL_PATH` points at a local Diffusers checkpoint directory.\n",
" \"\"\"\n",
" override = os.environ.get(\"COSMOS3_FP8_MODEL_PATH\")\n",
" if override:\n",
" return Path(override).expanduser().resolve()\n",
" return resolve_model_id(model)\n",
"\n",
"\n",
"def resolve_model_id(model: str) -> str:\n",
" return MODEL_IDS.get(model, model)\n",
"\n",
Expand Down Expand Up @@ -680,37 +715,60 @@
"def get_pipe(model: str):\n",
" global _pipe, _pipe_model\n",
" model_id = resolve_model_id(model)\n",
" if _pipe is not None and _pipe_model == model_id:\n",
" prec = precision_mode()\n",
" cache_key = f\"{model_id}@{prec}\"\n",
" if _pipe is not None and _pipe_model == cache_key:\n",
" return _pipe\n",
" release_pipe()\n",
" diffusers_logging.set_verbosity_info()\n",
" print(f\"loading {model_id}...\")\n",
" t0 = time.time()\n",
" fp8 = is_fp8_precision()\n",
" if fp8 and model not in FP8_SUPPORTED_MODELS:\n",
" raise RuntimeError(\n",
" f\"FP8 is only supported for {sorted(FP8_SUPPORTED_MODELS)}; got {model}. \"\n",
" \"Unset COSMOS3_DIFFUSERS_PRECISION or set it to bf16.\"\n",
" )\n",
" if fp8:\n",
" # ModelOpt FP8 graph restore + kernels (required before from_pretrained).\n",
" import modelopt.torch.quantization.backends.fp8_per_tensor_gemm # noqa: F401\n",
" from modelopt.torch.opt import enable_huggingface_checkpointing\n",
"\n",
" enable_huggingface_checkpointing()\n",
" pretrained_source = resolve_fp8_pretrained_source(model)\n",
" from_kwargs = {\"token\": os.environ.get(\"HF_TOKEN\") or None}\n",
" if not os.environ.get(\"COSMOS3_FP8_MODEL_PATH\"):\n",
" from_kwargs[\"revision\"] = os.environ.get(\"COSMOS3_FP8_REVISION\", \"fp8\")\n",
" print(f\"loading FP8 {pretrained_source}...\")\n",
" else:\n",
" pretrained_source = model_id\n",
" from_kwargs = {\"token\": os.environ.get(\"HF_TOKEN\") or None}\n",
" print(f\"loading {model_id}...\")\n",
"\n",
" if is_distilled_model(model):\n",
" if Cosmos3DistilledModularPipeline is None:\n",
" raise RuntimeError(\n",
" \"This diffusers build does not provide Cosmos3DistilledModularPipeline. \"\n",
" \"Reinstall diffusers from a revision that includes the Cosmos3 distilled pipeline.\"\n",
" )\n",
" pipe = Cosmos3DistilledModularPipeline.from_pretrained(\n",
" model_id,\n",
" token=os.environ.get(\"HF_TOKEN\") or None,\n",
" pretrained_source,\n",
" **from_kwargs,\n",
" )\n",
" pipe.load_components(torch_dtype=torch.bfloat16)\n",

@yzhautouskay yzhautouskay Aug 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure what the final model-card/repository structure will be, but we need important fix for the current nvidia/Cosmos3-Experimental subfolders to fix Distilled FP8. Because Distilled models can only be used with Modular pipeline, and how weights and quantization state get loaded for regular and modular pipelines differ.

Regular (called internally by Diffusers):
  Cosmos3OmniTransformer.from_pretrained(
      ".../cosmos3-nano-fp8-14072026/transformer"
  )

  ModelOpt finds:
  ".../cosmos3-nano-fp8-14072026/transformer/modelopt_state.pth"  ✓

Distilled (called internally by load_components):
  Cosmos3OmniTransformer.from_pretrained(
      "nvidia/Cosmos3-Experimental",
      revision="refs/pr/17",
      subfolder="cosmos3-super-t2i-4step-fp8-14072026/transformer",
  )

  ModelOpt checks:
  "nvidia/Cosmos3-Experimental/modelopt_state.pth"  ✗

  Actual state:
  "nvidia/Cosmos3-Experimental/
   cosmos3-super-t2i-4step-fp8-14072026/transformer/modelopt_state.pth"

The regular loader joins the component path first, while the modular loader keeps the subfolder separate, which ModelOpt ignores during state lookup.

Fixing it here is safer than changing generic diffusers loading or depending fixing in ModelOpt, which will only land in 0.46+ version and will be incompatible with 0.44 used for quantization already.
So the fix is -- load the transformer from its full path, pin it, then load the remaining components normally:

Suggested change
" pipe.load_components(torch_dtype=torch.bfloat16)\n",
" if fp8:\n",
" # Pre-load the transformer from its full path so ModelOpt 0.44 finds\n",
" # transformer/modelopt_state.pth and restores the FP8 quantizer graph.\n",
" from diffusers import AutoModel\n",
"\n",
" transformer_dir = Path(pretrained_source) / \"transformer\"\n",
" if not (transformer_dir / \"modelopt_state.pth\").is_file():\n",
" raise RuntimeError(\n",
" f\"FP8 requested but no ModelOpt state at {transformer_dir / 'modelopt_state.pth'}. \"\n",
" \"Point COSMOS3_FP8_MODEL_PATH at a pre-quantized distilled export directory.\"\n",
" )\n",
" transformer = AutoModel.from_pretrained(str(transformer_dir), torch_dtype=torch.bfloat16)\n",
" pipe.update_components(transformer=transformer)\n",
" pipe.load_components(\n",
" pretrained_model_name_or_path=distilled_components_root(pretrained_source),\n",
" torch_dtype=torch.bfloat16,\n",
" )\n",
" else:\n",
" pipe.load_components(torch_dtype=torch.bfloat16)\n",

@yzhautouskay yzhautouskay Aug 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There might be a better solution, so feel free to handle it other way. This is also something to fix for the HF model cards examples, although this is dependent on the final HF repo structure.

Some helper that I used to verify correctness:

def verify_fp8(model: str) -> None:
    transformer = get_pipe(model).transformer
    tensors = list(transformer.named_parameters()) + list(transformer.named_buffers())

    fp8 = sum(p.dtype == torch.float8_e4m3fn for p in transformer.parameters())
    quantizers = sum("quantizer" in name.lower() for name, _ in transformer.named_modules())
    meta = sum(t.is_meta for _, t in tensors)

    # Float8 weights alone are insufficient; quantizers prove ModelOpt restored the scales.
    if not fp8 or not quantizers or meta:
        raise RuntimeError(f"Invalid FP8 load: weights={fp8}, quantizers={quantizers}, meta={meta}")

    print(f"FP8 verified: weights={fp8}, quantizers={quantizers}, meta=0")

@ConstBob ConstBob Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the review and suggestions! I implemented the Distilled FP8 preload approach you suggested (full-path transformer load → update_components → load_components, plus a small helper to materialize Hub checkpoints locally) and added your verify_fp8 helper: f4878b6
Below testing has been conducted: On A100-80GB cluster (nvidia-modelopt==0.44.0), I loaded Experimental Distilled T2I 4-Step FP8 through that path and ran verify_fp8("Cosmos3-Super-Text2Image-4Step"); ModelOpt restored transformer/modelopt_state.pth and verification passed (weights=896, quantizers=2709, meta=0).

@yzhautouskay yzhautouskay Aug 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! Let's wait till checkpoints are released in the final HF model repos to update paths in the cookbook before merging it

" pipe.enable_safety_checker()\n",
" else:\n",
" pipe = Cosmos3OmniPipeline.from_pretrained(\n",
" model_id,\n",
" pretrained_source,\n",
" torch_dtype=torch.bfloat16,\n",
" safety_checker=None,\n",
" enable_safety_checker=True,\n",
" token=os.environ.get(\"HF_TOKEN\") or None,\n",
" **from_kwargs,\n",
" )\n",
" pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=FIXED_SAMPLING[\"shift\"])\n",
" pipe.to(\"cuda\")\n",
" _pipe = pipe\n",
" _pipe_model = model_id\n",
" print(f\"loaded pipeline in {time.time() - t0:.1f}s\")\n",
" _pipe_model = cache_key\n",
" print(f\"loaded pipeline in {time.time() - t0:.1f}s (precision={prec})\")\n",
" return _pipe\n",
"\n",
"\n",
Expand Down Expand Up @@ -839,10 +897,45 @@
" for src in images:\n",
" print(f\"source: {src} ({src.stat().st_size // 1024} KB)\")\n",
" display(Image(filename=str(src), width=720))\n"
],
]
},
{
"cell_type": "markdown",
"id": "quantized-checkpoints-fp8",
"metadata": {},
"source": [
"## Quantized Checkpoints (FP8)\n",
"\n",
"FP8 builds are available for **`Cosmos3-Nano`**, **`Cosmos3-Super`**, and the Super **4-Step** distilled checkpoints (`Cosmos3-Edge` has no FP8 build). Serving FP8 is a checkpoint-loading change: the use-case request cells below stay the same.\n",
"\n",
"Static-scale FP8 checkpoints are produced with [NVIDIA TensorRT Model Optimizer](https://github.com/NVIDIA/TensorRT-Model-Optimizer) and load through the same Diffusers pipelines as bf16 (`Cosmos3OmniPipeline` for Nano/Super, `Cosmos3DistilledModularPipeline` for 4-Step). The install cell pins `nvidia-modelopt==0.44.0`.\n",
"\n",
"Enable FP8 before running the helper cell (or re-run the helper cell after changing the env):\n",
"\n",
"```bash\n",
"export COSMOS3_DIFFUSERS_PRECISION=fp8\n",
"```\n",
"\n",
"With FP8 enabled, this notebook loads the public model id from `MODEL_IDS` with `revision=fp8` (override with `COSMOS3_FP8_REVISION`). Optionally set `COSMOS3_FP8_MODEL_PATH` to a local Diffusers checkpoint directory instead.\n",
"\n",
"Loading may print a `Some weights of the model checkpoint ... not used` warning for `weight_scale` / `input_scale` keys. Those are vLLM-Omni scale tensors shipped alongside the Diffusers ModelOpt state; the warning is expected.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "quantized-checkpoints-fp8-enable",
"metadata": {},
"outputs": [],
"id": "77e18231"
"source": [
"# Optional: flip precision for subsequent get_pipe / run_diffusers_payload calls.\n",
"# Re-run the helper cell above after changing this if helpers were already executed.\n",
"import os\n",
"\n",
"# os.environ[\"COSMOS3_DIFFUSERS_PRECISION\"] = \"fp8\" # uncomment to use FP8 Nano/Super/4-Step\n",
"print(\"COSMOS3_DIFFUSERS_PRECISION:\", os.environ.get(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\"))\n",
"print(\"COSMOS3_FP8_REVISION:\", os.environ.get(\"COSMOS3_FP8_REVISION\", \"fp8\"))\n"
]
},
{
"cell_type": "markdown",
Expand Down