From 856a1a2d5bb30a34c4e8e3c263237ec22cc0fa9d Mon Sep 17 00:00:00 2001 From: Yiming Zhao Date: Wed, 21 Oct 2026 11:14:39 -0700 Subject: [PATCH 1/3] cookbook(diffusers): restore FP8 checkpoints from a fresh Hub cache Port the Diffusers FP8 loading path onto the cookbooks now under .agents/skills, including ModelOpt state prefetch and quantizer validation. --- .../action/run_fd_with_diffusers.ipynb | 145 ++++++++++-- .../action/run_id_with_diffusers.ipynb | 147 ++++++++++-- .../audiovisual/run_with_diffusers.ipynb | 210 ++++++++++++++++-- .../run_video_transfer_with_diffusers.ipynb | 156 +++++++++++-- 4 files changed, 604 insertions(+), 54 deletions(-) diff --git a/.agents/skills/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb b/.agents/skills/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb index d0cc9ab05..712c871cc 100644 --- a/.agents/skills/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb +++ b/.agents/skills/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb @@ -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" ] }, { @@ -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" ] }, { @@ -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", @@ -126,6 +130,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:\", \"\" if os.environ.get(\"HF_TOKEN\") else \"\")\n", @@ -133,7 +138,7 @@ " 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" ] }, { @@ -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", @@ -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" ] }, { @@ -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", @@ -274,6 +281,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:\", \"\" if os.environ.get(\"HF_TOKEN\") else \"\")\n", @@ -281,7 +289,7 @@ " 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" ] }, { @@ -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", @@ -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", diff --git a/.agents/skills/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb b/.agents/skills/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb index e32ee1e6d..e3e4bdafe 100644 --- a/.agents/skills/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb +++ b/.agents/skills/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb @@ -18,7 +18,9 @@ "\n", "Inverse dynamics runs the other direction from forward dynamics: the input is a clip and no trajectory, and the model recovers the motion that connects the frames. It runs `Cosmos3OmniPipeline` with the clip grouped into a `CosmosActionCondition`.\n", "\n", - "It runs the two AV clips plus the Bridge, AgiBotWorld-Beta, RoboMIND Franka, RoboMIND Franka dual-arm, RoboMIND UR, UMI, and Fractal episodes under [`assets/`](./assets). For forward dynamics, see [`run_fd_with_diffusers.ipynb`](./run_fd_with_diffusers.ipynb)." + "It runs the two AV clips plus the Bridge, AgiBotWorld-Beta, RoboMIND Franka, RoboMIND Franka dual-arm, RoboMIND UR, UMI, and Fractal episodes under [`assets/`](./assets). For forward dynamics, see [`run_fd_with_diffusers.ipynb`](./run_fd_with_diffusers.ipynb).\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" ] }, { @@ -61,7 +63,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" ] }, { @@ -112,6 +115,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", @@ -124,6 +128,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:\", \"\" if os.environ.get(\"HF_TOKEN\") else \"\")\n", @@ -131,7 +136,7 @@ " 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" ] }, { @@ -183,7 +188,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", @@ -195,7 +201,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" ] }, { @@ -260,6 +266,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", @@ -272,6 +279,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:\", \"\" if os.environ.get(\"HF_TOKEN\") else \"\")\n", @@ -279,7 +287,7 @@ " 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" ] }, { @@ -374,7 +382,7 @@ "\n", "Inverse dynamics reads an existing clip and predicts the action chunk that connects its frames, so it conditions on `action.video` and passes no `raw_actions`. The only output is the predicted action array; nothing is generated to view as a video.\n", "\n", - "Predicted AV actions are 9D framewise deltas \u2014 a 3D translation plus a 6D rotation \u2014 which integrate into absolute camera poses for plotting.\n", + "Predicted AV actions are 9D framewise deltas — a 3D translation plus a 6D rotation — which integrate into absolute camera poses for plotting.\n", "\n", "Robot episodes work the same way with two differences: their actions come back normalized, so they are denormalized with the episode's own statistics first, and a single action vector can carry several end-effectors, so each case names the 9D block to integrate. Longer episodes are covered by predicting one chunk per sampled window and concatenating the results." ] @@ -676,25 +684,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", @@ -915,6 +1005,35 @@ " )\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-usecases", diff --git a/.agents/skills/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb b/.agents/skills/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb index ae7199cd8..19aba75c9 100644 --- a/.agents/skills/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb +++ b/.agents/skills/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb @@ -19,7 +19,9 @@ "\n", "Run the Cosmos3-Nano, Cosmos3-Super, or Cosmos3-Edge examples independently. Cosmos3-Edge has no audio modules and uses its 480p generation settings. The notebook also includes Nano and Super video-to-video examples, while Super T2V/I2V examples run without audio. Each section loads the matching model explicitly.\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" + "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", + "\n", + "FP8 checkpoints: set `COSMOS3_DIFFUSERS_PRECISION=fp8` before loading. Nano, Super, and Edge use the dedicated Hub repos `nvidia/Cosmos3-Nano-FP8`, `nvidia/Cosmos3-Super-FP8`, and `nvidia/Cosmos3-Edge-FP8`. Super 4-Step T2I/I2V keep the existing model cards with `revision=fp8`. Use-case cells stay the same.\n" ] }, { @@ -57,6 +59,7 @@ "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" ] }, @@ -103,6 +106,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", @@ -114,6 +118,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:\", \"\" if os.environ.get(\"HF_TOKEN\") else \"\")\n", @@ -159,7 +164,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", @@ -596,7 +602,6 @@ " else:\n", " print(f\"image: {vision_display_path.relative_to(COSMOS_ROOT)}\")\n", " display(Image(filename=str(vision_display_path), width=420))\n", - " print(json.dumps({k: payload[k] for k in [\"model_mode\", \"name\", \"enable_sound\", \"num_steps\", \"guidance\", \"shift\", \"fps\", \"num_frames\", \"resolution\", \"aspect_ratio\", \"seed\"]}, indent=2))\n", " summary_keys = [\"model_mode\", \"name\", \"enable_sound\", \"num_steps\", \"guidance\", \"shift\", \"fps\", \"num_frames\", \"resolution\", \"aspect_ratio\", \"seed\"]\n", " print(json.dumps({k: payload[k] for k in summary_keys if k in payload}, indent=2))\n", " if payload.get(\"distilled\"):\n", @@ -644,18 +649,119 @@ " \"Cosmos3-Super-Image2Video-4Step\",\n", "}\n", "\n", + "# Base models ship FP8 on dedicated Hub repos. Distilled 4-Step FP8 stays on revision=fp8\n", + "# of the existing cards.\n", + "FP8_SEPARATE_REPOS = {\n", + " \"Cosmos3-Nano\": \"nvidia/Cosmos3-Nano-FP8\",\n", + " \"Cosmos3-Super\": \"nvidia/Cosmos3-Super-FP8\",\n", + " \"Cosmos3-Edge\": \"nvidia/Cosmos3-Edge-FP8\",\n", + "}\n", + "\n", + "FP8_SUPPORTED_MODELS = set(FP8_SEPARATE_REPOS) | set(DISTILLED_MODELS)\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_model_id(model: str) -> str:\n", " return MODEL_IDS.get(model, model)\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() and model in FP8_SEPARATE_REPOS:\n", + " return FP8_SEPARATE_REPOS[model]\n", + " return resolve_model_id(model)\n", + "\n", + "\n", + "def pretrained_from_kwargs(model: str) -> dict:\n", + " kwargs = {\"token\": os.environ.get(\"HF_TOKEN\") or None}\n", + " if (\n", + " is_fp8_precision()\n", + " and not os.environ.get(\"COSMOS3_FP8_MODEL_PATH\")\n", + " and model in DISTILLED_MODELS\n", + " ):\n", + " kwargs[\"revision\"] = os.environ.get(\"COSMOS3_FP8_REVISION\", \"fp8\")\n", + " return kwargs\n", + "\n", + "\n", "def is_distilled_model(model: str) -> bool:\n", " 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", + "\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", + "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", + " fp8, quantizers, meta = validate_fp8_transformer(get_pipe(model).transformer)\n", + " print(f\"FP8 verified: weights={fp8}, quantizers={quantizers}, meta={meta}\")\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", @@ -692,13 +798,31 @@ "\n", "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", + " 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", + " 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", + " from_kwargs = pretrained_from_kwargs(model)\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", + " print(f\"loading FP8 {pretrained_source}...\")\n", + " else:\n", + " print(f\"loading {pretrained_source}...\")\n", + "\n", " if is_distilled_model(model):\n", " if Cosmos3DistilledModularPipeline is None:\n", " raise RuntimeError(\n", @@ -706,24 +830,54 @@ " \"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", + " if fp8:\n", + " state_path = ensure_modelopt_state(\n", + " pretrained_source,\n", + " token=from_kwargs.get(\"token\"),\n", + " revision=from_kwargs.get(\"revision\"),\n", + " )\n", + " print(f\"ModelOpt state: {state_path}\")\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", + " 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\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", @@ -859,6 +1013,36 @@ " display(Image(filename=str(src), width=720))\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 / Super / Edge:** load `nvidia/Cosmos3-Nano-FP8`, `nvidia/Cosmos3-Super-FP8`, and `nvidia/Cosmos3-Edge-FP8` (default revision on those repos). Do not pass `revision=fp8` on the BF16 cards.\n", + "- **Super 4-Step T2I / I2V:** load the existing cards with `revision=fp8`.\n", + "\n", + "Diffusers follows the checkpoint schedule (W8A16 on the first and last few steps, W8A8 in between). Distilled 4-Step stays all W8A8. 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. Distilled FP8 preloads `transformer/` so ModelOpt 0.44 restores `modelopt_state.pth`. After load, `verify_fp8(\"Cosmos3-Nano\")` (or another supported model) checks FP8 weights and quantizer modules.\n", + "\n" + ] + }, + { + "cell_type": "code", + "id": "fp8-precision-toggle", + "metadata": {}, + "outputs": [], + "execution_count": null, + "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": "d988d45b", diff --git a/.agents/skills/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb b/.agents/skills/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb index 57fe9f129..90dfa7069 100644 --- a/.agents/skills/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb +++ b/.agents/skills/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb @@ -18,7 +18,9 @@ "\n", "This notebook runs Cosmos3 structural-control transfer with `Cosmos3OmniModularPipeline`, the pipeline that exposes control inputs.\n", "\n", - "It runs edge, blur, depth, segmentation, and world-scenario-map transfer on Cosmos3-Nano, then the same five controls on Cosmos3-Super." + "It runs edge, blur, depth, segmentation, and world-scenario-map transfer on Cosmos3-Nano, then the same five controls on Cosmos3-Super.\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" ] }, { @@ -58,7 +60,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" ] }, { @@ -108,6 +111,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", @@ -120,12 +124,13 @@ " \"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:\", \"\" if os.environ.get(\"HF_TOKEN\") else \"\")\n", "\n", "\n", - "configure_diffusers_environment()" + "configure_diffusers_environment()\n" ] }, { @@ -165,7 +170,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", @@ -174,7 +180,7 @@ "echo\n", "echo \"Installed dependencies into: $COSMOS3_DIFFUSERS_VENV\"\n", "echo \"Next: switch this notebook kernel to: Cosmos3 Diffusers (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" ] }, { @@ -238,6 +244,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", @@ -250,12 +257,13 @@ " \"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:\", \"\" if os.environ.get(\"HF_TOKEN\") else \"\")\n", "\n", "\n", - "configure_diffusers_environment()" + "configure_diffusers_environment()\n" ] }, { @@ -435,19 +443,110 @@ " 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", + "\n", + "def materialize_checkpoint_dir(pretrained_source, token) -> 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(snapshot_download(str(pretrained_source), token=token)).resolve()\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) -> Cosmos3OmniModularPipeline:\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", - " print(f\"loading {model_id}...\")\n", " t0 = time.time()\n", - " pipe = Cosmos3OmniModularPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)\n", - " pipe.load_components(torch_dtype=torch.bfloat16)\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", + " print(f\"loading FP8 {pretrained_source}...\")\n", + " else:\n", + " print(f\"loading {pretrained_source}...\")\n", + " pipe = Cosmos3OmniModularPipeline.from_pretrained(\n", + " pretrained_source,\n", + " torch_dtype=torch.bfloat16,\n", + " token=token,\n", + " )\n", + " if fp8:\n", + " # Pre-load the transformer from its full path so ModelOpt 0.44 finds\n", + " # transformer/modelopt_state.pth. Modular load_components keeps subfolder\n", + " # separate, which ModelOpt ignores.\n", + " from diffusers import AutoModel\n", + "\n", + " ckpt_dir = materialize_checkpoint_dir(pretrained_source, token)\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 checkpoint directory.\"\n", + " )\n", + " transformer = AutoModel.from_pretrained(str(transformer_dir), torch_dtype=torch.bfloat16)\n", + " pipe.update_components(transformer=transformer)\n", + " # Keep the remaining component sources from modular_model_index.json. The FP8 repo's\n", + " # root modelopt_state.pth belongs to the Reasoner and must not be applied to the VAE or\n", + " # audio tokenizer.\n", + " pipe.load_components(torch_dtype=torch.bfloat16, token=token)\n", + " else:\n", + " pipe.load_components(torch_dtype=torch.bfloat16)\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", @@ -496,7 +595,36 @@ " print(f\"generated in {time.time() - t0:.1f}s\")\n", " export_to_video(videos, str(output_path), fps=spec[\"fps\"], macro_block_size=1)\n", " print(f\"wrote {output_path}\")\n", - " return output_path" + " 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. Transfer preloads `transformer/` so ModelOpt 0.44 restores `modelopt_state.pth`.\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" ] }, { From 9956d477cd9c6f5e519f2ff2d48f77fe0722eb58 Mon Sep 17 00:00:00 2001 From: Yiming Zhao Date: Mon, 26 Oct 2026 09:36:21 -0700 Subject: [PATCH 2/3] cookbook(diffusers): reject unsupported FP8 before releasing a pipeline Move the audiovisual support check ahead of release_pipe, and describe the FP8 weight check as a restore smoke check rather than GEMM verification. --- .../action/run_fd_with_diffusers.ipynb | 8 +++++-- .../action/run_id_with_diffusers.ipynb | 8 +++++-- .../audiovisual/run_with_diffusers.ipynb | 23 +++++++++++-------- .../run_video_transfer_with_diffusers.ipynb | 14 +++++++---- 4 files changed, 35 insertions(+), 18 deletions(-) diff --git a/.agents/skills/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb b/.agents/skills/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb index 712c871cc..374ef963d 100644 --- a/.agents/skills/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb +++ b/.agents/skills/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb @@ -634,7 +634,11 @@ "\n", "\n", "def validate_fp8_transformer(transformer) -> tuple[int, int, int]:\n", - " \"\"\"Reject FP8 weights loaded without their ModelOpt quantizer graph.\"\"\"\n", + " \"\"\"Smoke-check that ModelOpt restored FP8 weights and the quantizer graph.\n", + "\n", + " Counts float8 weights, quantizer modules, and meta tensors. It does not check\n", + " that the W8A8 GEMM path is active.\n", + " \"\"\"\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", @@ -676,7 +680,7 @@ " )\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", + " print(f\"FP8 state restored: weights={fp8_weights}, quantizers={quantizers}, meta={meta}\")\n", " pipe.to(\"cuda\")\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", diff --git a/.agents/skills/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb b/.agents/skills/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb index e3e4bdafe..76ab923ef 100644 --- a/.agents/skills/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb +++ b/.agents/skills/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb @@ -739,7 +739,11 @@ "\n", "\n", "def validate_fp8_transformer(transformer) -> tuple[int, int, int]:\n", - " \"\"\"Reject FP8 weights loaded without their ModelOpt quantizer graph.\"\"\"\n", + " \"\"\"Smoke-check that ModelOpt restored FP8 weights and the quantizer graph.\n", + "\n", + " Counts float8 weights, quantizer modules, and meta tensors. It does not check\n", + " that the W8A8 GEMM path is active.\n", + " \"\"\"\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", @@ -781,7 +785,7 @@ " )\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", + " print(f\"FP8 state restored: weights={fp8_weights}, quantizers={quantizers}, meta={meta}\")\n", " pipe.to(\"cuda\")\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", diff --git a/.agents/skills/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb b/.agents/skills/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb index 19aba75c9..9ecae70e6 100644 --- a/.agents/skills/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb +++ b/.agents/skills/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb @@ -677,6 +677,11 @@ "\n", "def resolve_pretrained_source(model: str):\n", " \"\"\"Hub repo id or local path for the requested precision.\"\"\"\n", + " if is_fp8_precision() 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", " override = os.environ.get(\"COSMOS3_FP8_MODEL_PATH\")\n", " if is_fp8_precision() and override:\n", " return Path(override).expanduser().resolve()\n", @@ -742,7 +747,11 @@ "\n", "\n", "def validate_fp8_transformer(transformer) -> tuple[int, int, int]:\n", - " \"\"\"Reject FP8 weights loaded without their ModelOpt quantizer graph.\"\"\"\n", + " \"\"\"Smoke-check that ModelOpt restored FP8 weights and the quantizer graph.\n", + "\n", + " Counts float8 weights, quantizer modules, and meta tensors. It does not check\n", + " that the W8A8 GEMM path is active.\n", + " \"\"\"\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", @@ -757,9 +766,9 @@ "\n", "\n", "def verify_fp8(model: str) -> None:\n", - " \"\"\"Confirm ModelOpt restored FP8 weights and quantizers (not just float8 tensors).\"\"\"\n", + " \"\"\"Confirm ModelOpt restored FP8 weights and quantizers, not that GEMM is active.\"\"\"\n", " fp8, quantizers, meta = validate_fp8_transformer(get_pipe(model).transformer)\n", - " print(f\"FP8 verified: weights={fp8}, quantizers={quantizers}, meta={meta}\")\n", + " print(f\"FP8 state restored: weights={fp8}, quantizers={quantizers}, meta={meta}\")\n", "\n", "\n", "def cuda_allocated_gib() -> float:\n", @@ -807,11 +816,6 @@ " diffusers_logging.set_verbosity_info()\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", " from_kwargs = pretrained_from_kwargs(model)\n", " if fp8:\n", " # ModelOpt FP8 graph restore + kernels (required before from_pretrained).\n", @@ -848,6 +852,7 @@ " )\n", " transformer = AutoModel.from_pretrained(str(transformer_dir), torch_dtype=torch.bfloat16)\n", " pipe.update_components(transformer=transformer)\n", + " # names=None loads only components that are still unset, so the preloaded transformer stays.\n", " pipe.load_components(\n", " pretrained_model_name_or_path=distilled_components_root(ckpt_dir),\n", " torch_dtype=torch.bfloat16,\n", @@ -873,7 +878,7 @@ " pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=FIXED_SAMPLING[\"shift\"])\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", + " print(f\"FP8 state restored: weights={fp8_weights}, quantizers={quantizers}, meta={meta}\")\n", " pipe.to(\"cuda\")\n", " _pipe = pipe\n", " _pipe_model = cache_key\n", diff --git a/.agents/skills/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb b/.agents/skills/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb index 90dfa7069..59a1e565c 100644 --- a/.agents/skills/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb +++ b/.agents/skills/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb @@ -485,7 +485,11 @@ "\n", "\n", "def validate_fp8_transformer(transformer) -> tuple[int, int, int]:\n", - " \"\"\"Reject FP8 weights loaded without their ModelOpt quantizer graph.\"\"\"\n", + " \"\"\"Smoke-check that ModelOpt restored FP8 weights and the quantizer graph.\n", + "\n", + " Counts float8 weights, quantizer modules, and meta tensors. It does not check\n", + " that the W8A8 GEMM path is active.\n", + " \"\"\"\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", @@ -535,15 +539,15 @@ " )\n", " transformer = AutoModel.from_pretrained(str(transformer_dir), torch_dtype=torch.bfloat16)\n", " pipe.update_components(transformer=transformer)\n", - " # Keep the remaining component sources from modular_model_index.json. The FP8 repo's\n", - " # root modelopt_state.pth belongs to the Reasoner and must not be applied to the VAE or\n", - " # audio tokenizer.\n", + " # names=None loads only components that are still unset, so the preloaded transformer stays.\n", + " # Remaining components keep the sources in modular_model_index.json. The FP8 repo's root\n", + " # modelopt_state.pth belongs to the Reasoner and must not be applied to the VAE or audio tokenizer.\n", " pipe.load_components(torch_dtype=torch.bfloat16, token=token)\n", " else:\n", " pipe.load_components(torch_dtype=torch.bfloat16)\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", + " print(f\"FP8 state restored: weights={fp8_weights}, quantizers={quantizers}, meta={meta}\")\n", " pipe.to(\"cuda\")\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", From 94a84508cebb302cb17030f1ff8217725f51a6c5 Mon Sep 17 00:00:00 2001 From: Yiming Zhao Date: Wed, 28 Oct 2026 16:07:44 -0700 Subject: [PATCH 3/3] cookbook(diffusers): include the FP8 revision in the pipeline cache key Distilled checkpoints share one repository, so a cache keyed only by repo and precision reused the previous pipeline after COSMOS3_FP8_REVISION changed. --- .../cosmos3/generator/audiovisual/run_with_diffusers.ipynb | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/.agents/skills/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb b/.agents/skills/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb index 9ecae70e6..c2d6596bc 100644 --- a/.agents/skills/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb +++ b/.agents/skills/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb @@ -809,14 +809,19 @@ " global _pipe, _pipe_model\n", " prec = precision_mode()\n", " pretrained_source = resolve_pretrained_source(model)\n", + " from_kwargs = pretrained_from_kwargs(model)\n", + " # Distilled FP8 keeps one repo id and selects the checkpoint with revision.\n", + " # The revision has to be in the cache key before the hit check; otherwise\n", + " # changing COSMOS3_FP8_REVISION returns the pipeline loaded for the previous revision.\n", " cache_key = f\"{pretrained_source}@{prec}\"\n", + " if \"revision\" in from_kwargs:\n", + " cache_key = f\"{cache_key}@{from_kwargs['revision']}\"\n", " if _pipe is not None and _pipe_model == cache_key:\n", " return _pipe\n", " release_pipe()\n", " diffusers_logging.set_verbosity_info()\n", " t0 = time.time()\n", " fp8 = is_fp8_precision()\n", - " from_kwargs = pretrained_from_kwargs(model)\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",