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 all 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
176 changes: 164 additions & 12 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 All @@ -643,6 +678,43 @@
" return model in DISTILLED_MODELS\n",
"\n",
"\n",
"def materialize_checkpoint_dir(pretrained_source, from_kwargs: dict) -> Path:\n",
" \"\"\"Return a local checkpoint root (download from Hub when given a repo id).\"\"\"\n",
" path = Path(str(pretrained_source))\n",
" if path.is_dir():\n",
" return path.resolve()\n",
" from huggingface_hub import snapshot_download\n",
"\n",
" return Path(\n",
" snapshot_download(\n",
" str(pretrained_source),\n",
" revision=from_kwargs.get(\"revision\"),\n",
" token=from_kwargs.get(\"token\"),\n",
" )\n",
" ).resolve()\n",
"\n",
"\n",
"def distilled_components_root(checkpoint_dir: Path) -> str:\n",
" \"\"\"Local Diffusers root for modular `load_components` after transformer preload.\"\"\"\n",
" return str(Path(checkpoint_dir).resolve())\n",
"\n",
"\n",
"def verify_fp8(model: str) -> None:\n",
" \"\"\"Confirm ModelOpt restored FP8 weights and quantizers (not just float8 tensors).\"\"\"\n",
" transformer = get_pipe(model).transformer\n",
" tensors = list(transformer.named_parameters()) + list(transformer.named_buffers())\n",
"\n",
" fp8 = sum(p.dtype == torch.float8_e4m3fn for p in transformer.parameters())\n",
" quantizers = sum(\"quantizer\" in name.lower() for name, _ in transformer.named_modules())\n",
" meta = sum(t.is_meta for _, t in tensors)\n",
"\n",
" # Float8 weights alone are insufficient; quantizers prove ModelOpt restored the scales.\n",
" if not fp8 or not quantizers or meta:\n",
" raise RuntimeError(f\"Invalid FP8 load: weights={fp8}, quantizers={quantizers}, meta={meta}\")\n",
"\n",
" print(f\"FP8 verified: weights={fp8}, quantizers={quantizers}, meta=0\")\n",
"\n",
"\n",
"def cuda_allocated_gib() -> float:\n",
" return torch.cuda.memory_allocated() / 1024**3 if torch.cuda.is_available() else 0.0\n",
"\n",
Expand Down Expand Up @@ -680,37 +752,80 @@
"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",
" 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",
" # Modular load_components uses repo+subfolder; ModelOpt ignores subfolder.\n",
" from diffusers import AutoModel\n",
"\n",
" ckpt_dir = materialize_checkpoint_dir(pretrained_source, from_kwargs)\n",
" transformer_dir = ckpt_dir / \"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(ckpt_dir),\n",
" torch_dtype=torch.bfloat16,\n",
" )\n",
" else:\n",
" pipe.load_components(torch_dtype=torch.bfloat16)\n",
" 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 +954,47 @@
" 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",
"For **4-Step distilled FP8**, ModelOpt 0.44 does not honor Hub `subfolder=` when restoring `modelopt_state.pth`. The helper therefore materializes a local checkpoint root, pre-loads `transformer/` via its full path, pins it with `update_components`, then loads the remaining modular components. After load you can sanity-check restore with `verify_fp8(\"Cosmos3-Super-Text2Image-4Step\")` (or the I2V 4-Step id).\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