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 11 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 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
619afb8
cookbook(diffusers): point FP8 loads at dedicated Hub repos
ConstBob Sep 15, 2026
a3bd513
Merge NVIDIA/cosmos main into feat/diffusers-fp8
ConstBob Sep 15, 2026
c4ecc2d
cookbook(diffusers): mention checkpoint mixed W8A8/W8A16 schedule
ConstBob Sep 17, 2026
86a24b4
Merge NVIDIA/cosmos main into feat/diffusers-fp8
ConstBob Sep 28, 2026
b8fdd60
cookbook(diffusers): load action and transfer FP8 from dedicated Hub …
ConstBob Sep 28, 2026
9ed508b
cookbook(diffusers): fix pose_utils import in ID notebook
ConstBob Sep 29, 2026
7f7db6a
Merge remote-tracking branch 'upstream/main' into feat/diffusers-fp8
ConstBob Oct 19, 2026
dd7088e
cookbook(diffusers): restore FP8 state from fresh Hub cache
ConstBob Oct 19, 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
145 changes: 132 additions & 13 deletions cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@
"\n",
"It runs three AV ego trajectories, three camera-pose trajectories, an autoregressive multiview DROID rollout, an autoregressive UMI rollout, and a bimanual hand-pose chunk. For inverse dynamics, see [`run_id_with_diffusers.ipynb`](./run_id_with_diffusers.ipynb).\n",
"\n",
"> **Guardrails:** Disabled by default in this notebook. Set `COSMOS3_DIFFUSERS_GUARDRAILS=true` before running the helper cell to enable them."
"> **Guardrails:** Disabled by default in this notebook. Set `COSMOS3_DIFFUSERS_GUARDRAILS=true` before running the helper cell to enable them.\n",
"\n",
"FP8 checkpoints: set `COSMOS3_DIFFUSERS_PRECISION=fp8` before loading. Nano and Super use the dedicated Hub repos `nvidia/Cosmos3-Nano-FP8` and `nvidia/Cosmos3-Super-FP8`. Use-case cells stay the same.\n"
]
},
{
Expand Down Expand Up @@ -63,7 +65,8 @@
"export HF_HOME=/path/to/large/huggingface/cache\n",
"export UV_LINK_MODE=copy\n",
"export CUDA_VISIBLE_DEVICES=0\n",
"```"
"export COSMOS3_DIFFUSERS_PRECISION=bf16 # or fp8\n",
"```\n"
]
},
{
Expand Down Expand Up @@ -114,6 +117,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 @@ -126,14 +130,15 @@
" \"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",
" if not (COSMOS3_REPO / \"cosmos_framework\").is_dir():\n",
" print(f\"note: no cosmos_framework package under {COSMOS3_REPO}; point COSMOS3_REPO at your framework checkout\")\n",
"\n",
"\n",
"configure_diffusers_environment()"
"configure_diffusers_environment()\n"
]
},
{
Expand Down Expand Up @@ -185,7 +190,8 @@
" scipy \\\n",
" torch \\\n",
" torchvision \\\n",
" transformers\n",
" transformers \\\n",
" \"nvidia-modelopt==0.44.0\"\n",
"\n",
"# Video decoding uses the PyAV fallback defined in the helpers cell, avoiding TorchCodec ABI\n",
"# coupling with the selected PyTorch backend.\n",
Expand All @@ -197,7 +203,7 @@
"echo\n",
"echo \"Installed dependencies into: $COSMOS3_DIFFUSERS_ACTION_VENV\"\n",
"echo \"Next: switch this notebook kernel to: Cosmos3 Diffusers Action (Python 3.13)\"\n",
"echo \"After switching kernels, run the Restore Environment cell below, then continue with Verify.\""
"echo \"After switching kernels, run the Restore Environment cell below, then continue with Verify.\"\n"
]
},
{
Expand Down Expand Up @@ -262,6 +268,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 @@ -274,14 +281,15 @@
" \"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",
" if not (COSMOS3_REPO / \"cosmos_framework\").is_dir():\n",
" print(f\"note: no cosmos_framework package under {COSMOS3_REPO}; point COSMOS3_REPO at your framework checkout\")\n",
"\n",
"\n",
"configure_diffusers_environment()"
"configure_diffusers_environment()\n"
]
},
{
Expand Down Expand Up @@ -571,25 +579,107 @@
" print(f\"released previous pipeline; cuda allocated {cuda_allocated_gib():.1f} GiB\")\n",
"\n",
"\n",
"\n",
"FP8_SEPARATE_REPOS = {\n",
" \"Cosmos3-Nano\": \"nvidia/Cosmos3-Nano-FP8\",\n",
" \"Cosmos3-Super\": \"nvidia/Cosmos3-Super-FP8\",\n",
"}\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_pretrained_source(model: str):\n",
" \"\"\"Hub repo id or local path for the requested precision.\"\"\"\n",
" override = os.environ.get(\"COSMOS3_FP8_MODEL_PATH\")\n",
" if is_fp8_precision() and override:\n",
" return Path(override).expanduser().resolve()\n",
" if is_fp8_precision():\n",
" if model not in FP8_SEPARATE_REPOS:\n",
" raise RuntimeError(\n",
" f\"FP8 is only supported for {sorted(FP8_SEPARATE_REPOS)}; got {model}. \"\n",
" \"Unset COSMOS3_DIFFUSERS_PRECISION or set it to bf16.\"\n",
" )\n",
" return FP8_SEPARATE_REPOS[model]\n",
" return MODEL_IDS.get(model, model)\n",
"\n",
"\n",
"def ensure_modelopt_state(pretrained_source, *, token=None, revision=None) -> Path:\n",
" \"\"\"Ensure ModelOpt state is present before Diffusers resolves its Hub allowlist.\"\"\"\n",
" checkpoint_root = Path(str(pretrained_source))\n",
" if isinstance(pretrained_source, Path) or checkpoint_root.is_absolute():\n",
" state_path = checkpoint_root / \"transformer\" / \"modelopt_state.pth\"\n",
" else:\n",
" from huggingface_hub import hf_hub_download\n",
"\n",
" state_path = Path(\n",
" hf_hub_download(\n",
" repo_id=str(pretrained_source),\n",
" filename=\"transformer/modelopt_state.pth\",\n",
" revision=revision,\n",
" token=token,\n",
" )\n",
" )\n",
" if not state_path.is_file():\n",
" raise RuntimeError(\n",
" f\"FP8 requested but no ModelOpt state at {state_path}. \"\n",
" \"Use a ModelOpt FP8 checkpoint or set COSMOS3_FP8_MODEL_PATH to one.\"\n",
" )\n",
" return state_path\n",
"\n",
"\n",
"def validate_fp8_transformer(transformer) -> tuple[int, int, int]:\n",
" \"\"\"Reject FP8 weights loaded without their ModelOpt quantizer graph.\"\"\"\n",
" tensors = list(transformer.named_parameters()) + list(transformer.named_buffers())\n",
" fp8 = sum(parameter.dtype == torch.float8_e4m3fn for parameter in transformer.parameters())\n",
" quantizers = sum(\"quantizer\" in name.lower() for name, _ in transformer.named_modules())\n",
" meta = sum(tensor.is_meta for _, tensor in tensors)\n",
" if not fp8 or not quantizers or meta:\n",
" raise RuntimeError(f\"Invalid FP8 load: weights={fp8}, quantizers={quantizers}, meta={meta}\")\n",
" return fp8, quantizers, meta\n",
"\n",
"\n",
"def get_pipe(model: str) -> Cosmos3OmniPipeline:\n",
" global _pipe, _pipe_model\n",
" model_id = MODEL_IDS.get(model, model)\n",
" if _pipe is not None and _pipe_model == model_id:\n",
" prec = precision_mode()\n",
" pretrained_source = resolve_pretrained_source(model)\n",
" cache_key = f\"{pretrained_source}@{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",
" token = os.environ.get(\"HF_TOKEN\") or None\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",
" state_path = ensure_modelopt_state(pretrained_source, token=token)\n",
" print(f\"ModelOpt state: {state_path}\")\n",
" print(f\"loading FP8 {pretrained_source}...\")\n",
" else:\n",
" print(f\"loading {pretrained_source}...\")\n",
" pipe = Cosmos3OmniPipeline.from_pretrained(\n",
" model_id,\n",
" pretrained_source,\n",
" torch_dtype=torch.bfloat16,\n",
" safety_checker=None,\n",
" enable_safety_checker=GUARDRAILS,\n",
" token=os.environ.get(\"HF_TOKEN\") or None,\n",
" token=token,\n",
" )\n",
" if fp8:\n",
" fp8_weights, quantizers, meta = validate_fp8_transformer(pipe.transformer)\n",
" print(f\"FP8 verified: weights={fp8_weights}, quantizers={quantizers}, meta={meta}\")\n",
" pipe.to(\"cuda\")\n",
" _pipe, _pipe_model = pipe, model_id\n",
" print(f\"loaded pipeline in {time.time() - t0:.1f}s; cuda allocated {cuda_allocated_gib():.1f} GiB\")\n",
" _pipe, _pipe_model = pipe, cache_key\n",
" print(f\"loaded pipeline in {time.time() - t0:.1f}s (precision={prec}); cuda allocated {cuda_allocated_gib():.1f} GiB\")\n",
" return _pipe\n",
"\n",
"\n",
Expand Down Expand Up @@ -748,6 +838,35 @@
" return output_path\n"
]
},
{
"cell_type": "markdown",
"id": "fp8-quantized-checkpoints",
"metadata": {},
"source": [
"## Quantized Checkpoints (FP8)\n",
"\n",
"Default loads are BF16 (`COSMOS3_DIFFUSERS_PRECISION=bf16`). Set the env var to `fp8` before running the helper cell's `get_pipe` (or re-run that cell after changing it). Use-case request cells stay the same.\n",
"\n",
"Nano and Super load `nvidia/Cosmos3-Nano-FP8` and `nvidia/Cosmos3-Super-FP8`. Do not pass `revision=fp8` on the BF16 cards.\n",
"\n",
"Diffusers follows the checkpoint schedule (W8A16 on the first and last few steps, W8A8 in between). See the [Cosmos 3 Diffusers docs](https://huggingface.co/docs/diffusers/main/en/api/pipelines/cosmos3).\n",
"\n",
"Optional: point `COSMOS3_FP8_MODEL_PATH` at a local Diffusers checkpoint directory.\n",
"\n"
]
},
{
"cell_type": "code",
"id": "fp8-precision-toggle",
"metadata": {},
"execution_count": null,
"outputs": [],
"source": [
"# Default is bf16. Uncomment to load FP8 checkpoints in the cells below.\n",
"# os.environ[\"COSMOS3_DIFFUSERS_PRECISION\"] = \"fp8\"\n",
"print(\"COSMOS3_DIFFUSERS_PRECISION:\", os.environ.get(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\"))\n"
]
},
{
"cell_type": "markdown",
"id": "action-diffusers-helpers-15",
Expand Down
Loading