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 d0cc9ab0..374ef963 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,111 @@ " 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", + " \"\"\"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", + " 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 state restored: 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 +842,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 e32ee1e6..76ab923e 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,111 @@ " 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", + " \"\"\"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", + " 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 state restored: 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 +1009,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/transfer/run_video_transfer_with_diffusers.ipynb b/.agents/skills/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb index 57fe9f12..59a1e565 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,114 @@ " 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", + " \"\"\"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", + " 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", + " # 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 state restored: 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 +599,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" ] }, { diff --git a/.agents/skills/cosmos3-generate-media/SKILL.md b/.agents/skills/cosmos3-generate-media/SKILL.md index ca299da1..dc3fcf6a 100644 --- a/.agents/skills/cosmos3-generate-media/SKILL.md +++ b/.agents/skills/cosmos3-generate-media/SKILL.md @@ -121,6 +121,7 @@ uv pip install --no-config --python "$COSMOS3_DIFFUSERS_VENV" --torch-backend=cu "diffusers @ git+https://github.com/huggingface/diffusers.git" \ accelerate av cosmos_guardrail huggingface_hub imageio imageio-ffmpeg \ torch torchvision transformers \ + "nvidia-modelopt==0.44.0" \ && uv pip uninstall --no-config --python "$COSMOS3_DIFFUSERS_VENV" opencv-python \ && uv pip install --no-config --python "$COSMOS3_DIFFUSERS_VENV" opencv-python-headless==5.0.0.93 "$COSMOS3_DIFFUSERS_VENV/bin/python" -c "import cv2; from diffusers import Cosmos3OmniPipeline; print('pipeline import OK')" # must pass before any checkpoint download @@ -133,6 +134,8 @@ skill directory (the one with `assets/`); do not compose the pipeline call from Edge (`t2i_edge`, `t2v_edge_noaudio`): run § 8.4 of `references/diffusers.md` the same way — its setup block plus only the block of the requested mode, verbatim; § 8.3 has the README's Edge settings and the video-mode rules (`num_frames=189`, `fps=24`, `export_to_video`). +FP8: set `COSMOS3_DIFFUSERS_PRECISION=fp8` before `get_pipe`. The loader, cache key, and checkpoint map are in `references/diffusers.md` (Quantized Checkpoints). BF16 remains the default. + ## Variant 2 — cosmos_framework (torchrun, in-repo) Reference: [`references/cosmos-framework.md`](references/cosmos-framework.md) — diff --git a/.agents/skills/cosmos3-generate-media/references/diffusers.md b/.agents/skills/cosmos3-generate-media/references/diffusers.md index 325a1f2d..6d9be55f 100644 --- a/.agents/skills/cosmos3-generate-media/references/diffusers.md +++ b/.agents/skills/cosmos3-generate-media/references/diffusers.md @@ -67,6 +67,8 @@ This notebook runs Cosmos3 audiovisual generation directly with `Cosmos3OmniPipe 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. +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. + ## 1. Prerequisites - Machine, Hugging Face model access, `uvx hf@latest auth login` / `HF_TOKEN`: see payloads.md § 2.2 (`run_with_diffusers.ipynb` § 1 entry). @@ -83,6 +85,7 @@ export COSMOS3_TORCH_BACKEND=cu130 export HF_HOME=/path/to/large/huggingface/cache export UV_LINK_MODE=copy export CUDA_VISIBLE_DEVICES=0 +export COSMOS3_DIFFUSERS_PRECISION=bf16 # or fp8 ``` Path configuration (`COSMOS_ROOT`, `COSMOS3_AUDIOVISUAL_ROOT`, `OUT` = `COSMOS3_AUDIOVISUAL_OUTPUT_ROOT`): see payloads.md § 4.1, which replaces the notebook's `find_repo_root()` and output-root lines. The remainder of the notebook's configuration cell runs on top of it: @@ -108,6 +111,7 @@ def configure_diffusers_environment() -> None: os.environ.setdefault("HF_HOME", str(Path.home() / ".cache" / "huggingface")) os.environ.setdefault("HF_HUB_DISABLE_XET", "1") os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0") + os.environ.setdefault("COSMOS3_DIFFUSERS_PRECISION", "bf16") configure_diffusers_environment() @@ -115,12 +119,13 @@ configure_diffusers_environment() Backend-only variables the cell sets (the notebook's echo loop printed these together with `COSMOS_ROOT`, `COSMOS3_AUDIOVISUAL_OUTPUT_ROOT`, `UV_CACHE_DIR`, `HF_HOME`, `CUDA_VISIBLE_DEVICES` and `HF_TOKEN: `/``; those shared variables are in payloads.md § 4.1 and § 4.4): -| Variable | Default in the cell | Documented example (export block above) | -| ------------------------ | --------------------------------------------------------------------- | --------------------------------------- | -| `COSMOS3_DIFFUSERS_VENV` | `COSMOS_ROOT / ".venv-cosmos3-diffusers"`, resolved and exported back | `/path/to/.venv-cosmos3-diffusers` | -| `COSMOS3_TORCH_BACKEND` | `"cu130"`, exported back | `cu130` (`cu128` for CUDA 12.8) | -| `UV_LINK_MODE` | `"copy"` (`setdefault`) | `copy` | -| `HF_HUB_DISABLE_XET` | `"1"` (`setdefault`) | — | +| Variable | Default in the cell | Documented example (export block above) | +| ----------------------------- | --------------------------------------------------------------------- | --------------------------------------------- | +| `COSMOS3_DIFFUSERS_VENV` | `COSMOS_ROOT / ".venv-cosmos3-diffusers"`, resolved and exported back | `/path/to/.venv-cosmos3-diffusers` | +| `COSMOS3_TORCH_BACKEND` | `"cu130"`, exported back | `cu130` (`cu128` for CUDA 12.8) | +| `UV_LINK_MODE` | `"copy"` (`setdefault`) | `copy` | +| `HF_HUB_DISABLE_XET` | `"1"` (`setdefault`) | — | +| `COSMOS3_DIFFUSERS_PRECISION` | `"bf16"` (`setdefault`) | `bf16` (`fp8` selects a quantized checkpoint) | ## 3. Install Diffusers Dependencies @@ -147,6 +152,7 @@ uv pip install --no-config --python "$COSMOS3_DIFFUSERS_VENV" --torch-backend="$ torch \ torchvision \ transformers \ + "nvidia-modelopt==0.44.0" \ && uv pip uninstall --no-config --python "$COSMOS3_DIFFUSERS_VENV" opencv-python \ && uv pip install --no-config --python "$COSMOS3_DIFFUSERS_VENV" opencv-python-headless==5.0.0.93 @@ -224,6 +230,114 @@ _pipe_model = None # resolve_model_id and is_distilled_model: see payloads.md § 5 +FP8_SEPARATE_REPOS = { + "Cosmos3-Nano": "nvidia/Cosmos3-Nano-FP8", + "Cosmos3-Super": "nvidia/Cosmos3-Super-FP8", + "Cosmos3-Edge": "nvidia/Cosmos3-Edge-FP8", +} + +FP8_SUPPORTED_MODELS = set(FP8_SEPARATE_REPOS) | set(DISTILLED_MODELS) + + +def precision_mode() -> str: + return os.environ.get("COSMOS3_DIFFUSERS_PRECISION", "bf16").strip().lower() + + +def is_fp8_precision() -> bool: + return precision_mode() == "fp8" + + +def resolve_pretrained_source(model: str): + """Hub repo id or local path for the requested precision.""" + if is_fp8_precision() and model not in FP8_SUPPORTED_MODELS: + raise RuntimeError( + f"FP8 is only supported for {sorted(FP8_SUPPORTED_MODELS)}; got {model}. " + "Unset COSMOS3_DIFFUSERS_PRECISION or set it to bf16." + ) + override = os.environ.get("COSMOS3_FP8_MODEL_PATH") + if is_fp8_precision() and override: + return Path(override).expanduser().resolve() + if is_fp8_precision() and model in FP8_SEPARATE_REPOS: + return FP8_SEPARATE_REPOS[model] + return resolve_model_id(model) + + +def pretrained_from_kwargs(model: str) -> dict: + kwargs = {"token": os.environ.get("HF_TOKEN") or None} + if ( + is_fp8_precision() + and not os.environ.get("COSMOS3_FP8_MODEL_PATH") + and model in DISTILLED_MODELS + ): + kwargs["revision"] = os.environ.get("COSMOS3_FP8_REVISION", "fp8") + return kwargs + + +def materialize_checkpoint_dir(pretrained_source, from_kwargs: dict) -> Path: + """Return a local checkpoint root (download from Hub when given a repo id).""" + path = Path(str(pretrained_source)) + if path.is_dir(): + return path.resolve() + from huggingface_hub import snapshot_download + + return Path( + snapshot_download( + str(pretrained_source), + revision=from_kwargs.get("revision"), + token=from_kwargs.get("token"), + ) + ).resolve() + + + +def ensure_modelopt_state(pretrained_source, *, token=None, revision=None) -> Path: + """Ensure ModelOpt state is present before Diffusers resolves its Hub allowlist.""" + checkpoint_root = Path(str(pretrained_source)) + if isinstance(pretrained_source, Path) or checkpoint_root.is_absolute(): + state_path = checkpoint_root / "transformer" / "modelopt_state.pth" + else: + from huggingface_hub import hf_hub_download + + state_path = Path( + hf_hub_download( + repo_id=str(pretrained_source), + filename="transformer/modelopt_state.pth", + revision=revision, + token=token, + ) + ) + if not state_path.is_file(): + raise RuntimeError( + f"FP8 requested but no ModelOpt state at {state_path}. " + "Use a ModelOpt FP8 checkpoint or set COSMOS3_FP8_MODEL_PATH to one." + ) + return state_path + + +def validate_fp8_transformer(transformer) -> tuple[int, int, int]: + """Smoke-check that ModelOpt restored FP8 weights and the quantizer graph. + + Counts float8 weights, quantizer modules, and meta tensors. It does not check + that the W8A8 GEMM path is active. + """ + tensors = list(transformer.named_parameters()) + list(transformer.named_buffers()) + fp8 = sum(parameter.dtype == torch.float8_e4m3fn for parameter in transformer.parameters()) + quantizers = sum("quantizer" in name.lower() for name, _ in transformer.named_modules()) + meta = sum(tensor.is_meta for _, tensor in tensors) + if not fp8 or not quantizers or meta: + raise RuntimeError(f"Invalid FP8 load: weights={fp8}, quantizers={quantizers}, meta={meta}") + return fp8, quantizers, meta + +def distilled_components_root(checkpoint_dir: Path) -> str: + """Local Diffusers root for modular `load_components` after transformer preload.""" + return str(Path(checkpoint_dir).resolve()) + + +def verify_fp8(model: str) -> None: + """Confirm ModelOpt restored FP8 weights and quantizers, not that GEMM is active.""" + fp8, quantizers, meta = validate_fp8_transformer(get_pipe(model).transformer) + print(f"FP8 state restored: weights={fp8}, quantizers={quantizers}, meta={meta}") + def cuda_allocated_gib() -> float: return torch.cuda.memory_allocated() / 1024**3 if torch.cuda.is_available() else 0.0 @@ -261,13 +375,31 @@ def release_pipe() -> None: def get_pipe(model: str): global _pipe, _pipe_model - model_id = resolve_model_id(model) - if _pipe is not None and _pipe_model == model_id: + prec = precision_mode() + pretrained_source = resolve_pretrained_source(model) + from_kwargs = pretrained_from_kwargs(model) + # Distilled FP8 keeps one repo id and selects the checkpoint with revision. + # The revision has to be in the cache key before the hit check; otherwise + # changing COSMOS3_FP8_REVISION returns the pipeline loaded for the previous revision. + cache_key = f"{pretrained_source}@{prec}" + if "revision" in from_kwargs: + cache_key = f"{cache_key}@{from_kwargs['revision']}" + if _pipe is not None and _pipe_model == cache_key: return _pipe release_pipe() diffusers_logging.set_verbosity_info() - print(f"loading {model_id}...") t0 = time.time() + fp8 = is_fp8_precision() + if fp8: + # ModelOpt FP8 graph restore + kernels (required before from_pretrained). + import modelopt.torch.quantization.backends.fp8_per_tensor_gemm # noqa: F401 + from modelopt.torch.opt import enable_huggingface_checkpointing + + enable_huggingface_checkpointing() + print(f"loading FP8 {pretrained_source}...") + else: + print(f"loading {pretrained_source}...") + if is_distilled_model(model): if Cosmos3DistilledModularPipeline is None: raise RuntimeError( @@ -275,24 +407,55 @@ def get_pipe(model: str): "Reinstall diffusers from a revision that includes the Cosmos3 distilled pipeline." ) pipe = Cosmos3DistilledModularPipeline.from_pretrained( - model_id, - token=os.environ.get("HF_TOKEN") or None, + pretrained_source, + **from_kwargs, ) - pipe.load_components(torch_dtype=torch.bfloat16) + if fp8: + # Pre-load the transformer from its full path so ModelOpt 0.44 finds + # transformer/modelopt_state.pth and restores the FP8 quantizer graph. + # Modular load_components uses repo+subfolder; ModelOpt ignores subfolder. + from diffusers import AutoModel + + ckpt_dir = materialize_checkpoint_dir(pretrained_source, from_kwargs) + transformer_dir = ckpt_dir / "transformer" + if not (transformer_dir / "modelopt_state.pth").is_file(): + raise RuntimeError( + f"FP8 requested but no ModelOpt state at {transformer_dir / 'modelopt_state.pth'}. " + "Point COSMOS3_FP8_MODEL_PATH at a pre-quantized distilled export directory." + ) + transformer = AutoModel.from_pretrained(str(transformer_dir), torch_dtype=torch.bfloat16) + pipe.update_components(transformer=transformer) + # names=None loads only components that are still unset, so the preloaded transformer stays. + pipe.load_components( + pretrained_model_name_or_path=distilled_components_root(ckpt_dir), + torch_dtype=torch.bfloat16, + ) + else: + pipe.load_components(torch_dtype=torch.bfloat16) pipe.enable_safety_checker() else: + if fp8: + state_path = ensure_modelopt_state( + pretrained_source, + token=from_kwargs.get("token"), + revision=from_kwargs.get("revision"), + ) + print(f"ModelOpt state: {state_path}") pipe = Cosmos3OmniPipeline.from_pretrained( - model_id, + pretrained_source, torch_dtype=torch.bfloat16, safety_checker=None, enable_safety_checker=True, - token=os.environ.get("HF_TOKEN") or None, + **from_kwargs, ) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=FIXED_SAMPLING["shift"]) + if fp8: + fp8_weights, quantizers, meta = validate_fp8_transformer(pipe.transformer) + print(f"FP8 state restored: weights={fp8_weights}, quantizers={quantizers}, meta={meta}") pipe.to("cuda") _pipe = pipe - _pipe_model = model_id - print(f"loaded pipeline in {time.time() - t0:.1f}s") + _pipe_model = cache_key + print(f"loaded pipeline in {time.time() - t0:.1f}s (precision={prec})") return _pipe @@ -393,6 +556,23 @@ def run_diffusers_payload(payload_path: Path, output_dir: str | Path, *, model: Viewer helpers (`display_video`, `view_run`) are Jupyter-only and are not migrated; `view_run`'s media discovery rule is in payloads.md § 4.2 (`view_run` discovery row). How each payload key reaches the pipeline call: payloads.md § 10 (diffusers column). +## Quantized Checkpoints (FP8) + +Default loads are BF16 (`COSMOS3_DIFFUSERS_PRECISION=bf16`). Set the env var to `fp8` before calling `get_pipe` in § 7. Use-case requests stay the same. + +- **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. +- **Super 4-Step T2I / I2V:** load the existing cards with `revision=fp8`. `get_pipe` includes that revision in the pipeline cache key, so changing `COSMOS3_FP8_REVISION` does not reuse the pipeline loaded for the previous revision. + +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). + +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. It counts float8 weights, quantizer modules, and meta tensors. It does not check that the W8A8 GEMM path is active. + +```python +# Default is bf16. Uncomment to load FP8 checkpoints. +# os.environ["COSMOS3_DIFFUSERS_PRECISION"] = "fp8" +print("COSMOS3_DIFFUSERS_PRECISION:", os.environ.get("COSMOS3_DIFFUSERS_PRECISION", "bf16")) +``` + ## Use Cases Workflow wording and the three-step pattern: see payloads.md § 11 (`run_with_diffusers.ipynb` entry). Every section below the notebook's "Use Cases" heading ran the same cell triplet, "Create Payload" / "Run" / "View Results":