Skip to content

Add ModelOpt FP8 Cosmos Framework inference #301

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

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

Add ModelOpt FP8 Cosmos Framework inference #301

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7cc502e
docs: add ModelOpt FP8 Cosmos Framework inference
KrzysztofTomala Jul 31, 2026
200cf52
docs: clarify FP8 checkpoint downloads
KrzysztofTomala Aug 11, 2026
b1eff4d
docs: support FSDP in FP8 reasoner cookbook
KrzysztofTomala Aug 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 79 additions & 0 deletions cookbooks/cosmos3/generator/audiovisual/run_with_cosmos_framework.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -624,6 +624,85 @@
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Nano FP8: Text to Image\n",
"\n",
"FP8 checkpoints are published in the same Hugging Face repositories as the standard checkpoints, under the `fp8` revision rather than in separate FP8 repositories. This example downloads `nvidia/Cosmos3-Nano` at revision `fp8`. To use another released FP8 checkpoint, change `FP8_MODEL_ID` and `FP8_CHECKPOINT` below; keep `FP8_REVISION=fp8`.\n",
"\n",
"### Create Payload"
],
"id": "fe1b922e"
},
{
"cell_type": "code",
"metadata": {},
"source": [
"t2i_fp8_payload, t2i_fp8_output, t2i_fp8_model = create_payload(\"t2i\", backend=\"pytorch_fp8\")"
],
"execution_count": null,
"outputs": [],
"id": "73ff907c"
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Run"
],
"id": "aaf87ed7"
},
{
"cell_type": "code",
"metadata": {},
"source": [
"%%bash\n",
"set -euo pipefail\n",
"\n",
"cd \"$COSMOS3_REPO\"\n",
"FP8_MODEL_ID=\"nvidia/Cosmos3-Nano\"\n",
"FP8_REVISION=\"fp8\"\n",
"FP8_CHECKPOINT=\"$COSMOS3_REPO/checkpoints/Cosmos3-Nano-FP8\"\n",

@lfengad lfengad Aug 10, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

We might need to add some comments annotation for how to chang to other FP8 released checkpoints?

@KrzysztofTomala KrzysztofTomala Aug 11, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

I have added the comment

"uvx hf@latest download \"$FP8_MODEL_ID\" \\\n",
" --revision \"$FP8_REVISION\" \\\n",
" --local-dir \"$FP8_CHECKPOINT\"\n",
"\n",
"CUDA_VISIBLE_DEVICES=\"$CUDA_VISIBLE_DEVICES\" LD_LIBRARY_PATH= \\\n",
"\"$COSMOS3_UV_ENV/bin/torchrun\" \\\n",
" --nproc-per-node=\"$COSMOS3_NUM_GPUS\" \\\n",
" --master-addr=\"$COSMOS3_MASTER_ADDR\" \\\n",
" --master-port=\"$COSMOS3_TEXT_MASTER_PORT\" \\\n",
" -m cosmos_framework.scripts.inference \\\n",
" --parallelism-preset=throughput \\\n",
" -i \"$COSMOS3_PYTORCH_FP8_T2I_INPUT\" \\\n",
" -o \"$COSMOS3_PYTORCH_FP8_T2I_OUTPUT\" \\\n",
" --checkpoint-path \"$FP8_CHECKPOINT\" \\\n",
" --seed=0\n"
],
"execution_count": null,
"outputs": [],
"id": "e0cbe358"
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### View Results"
],
"id": "03e5c4b5"
},
{
"cell_type": "code",
"metadata": {},
"source": [
"view_run(t2i_fp8_output)"
],
"execution_count": null,
"outputs": [],
"id": "8970833a"
},
{
"cell_type": "markdown",
"metadata": {},
Expand Down
76 changes: 76 additions & 0 deletions cookbooks/cosmos3/reasoner/run_with_cosmos_framework.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -639,6 +639,82 @@
" print(json.dumps(json.loads(benchmark_path.read_text()).get(\"average\", {}), indent=2))\n"
]
},
{
"cell_type": "markdown",
"id": "da7c0e0a",
"metadata": {},
"source": [
"### 7a. Optional: Run Nano FP8 Text Inference\n",
"\n",
"FP8 checkpoints are published in the same Hugging Face repositories as the standard checkpoints, under the `fp8` revision rather than in separate FP8 repositories. This example downloads `nvidia/Cosmos3-Nano` at revision `fp8`. To use another released FP8 checkpoint, change `FP8_MODEL_ID` and `FP8_CHECKPOINT` below; keep `FP8_REVISION=fp8`.\n",
"\n",
"Set `FP8_NUM_GPUS=1` to run without FSDP. To shard the model with FSDP, expose multiple GPUs in `CUDA_VISIBLE_DEVICES` and set `FP8_NUM_GPUS` to their count (for example, `CUDA_VISIBLE_DEVICES=0,1` and `FP8_NUM_GPUS=2`).\n",
"\n",
"Expected output file:\n",
"\n",
"```text\n",
"packages/cosmos3/outputs/cookbooks/cosmos3/reasoner/nano/cosmos_framework_nano_fp8_text/nano_text/reasoner_text.txt\n",
"```\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "228e450b",
"metadata": {},
"outputs": [],
"source": [
"%%bash\n",
"set -euo pipefail\n",
"\n",
"cd \"$COSMOS3_REPO\"\n",
"FP8_MODEL_ID=\"nvidia/Cosmos3-Nano\"\n",
"FP8_REVISION=\"fp8\"\n",
"FP8_CHECKPOINT=\"$COSMOS3_REPO/checkpoints/Cosmos3-Nano-FP8\"\n",
"FP8_NUM_GPUS=\"${FP8_NUM_GPUS:-1}\"\n",
"uvx hf@latest download \"$FP8_MODEL_ID\" \\\n",
" --revision \"$FP8_REVISION\" \\\n",
" --local-dir \"$FP8_CHECKPOINT\"\n",
"\n",
"COSMOS_TRAINING=false CUDA_VISIBLE_DEVICES=\"$CUDA_VISIBLE_DEVICES\" \\\n",
".venv/bin/torchrun \\\n",
" --nproc-per-node=\"$FP8_NUM_GPUS\" \\\n",
" --master-addr=\"$COSMOS3_MASTER_ADDR\" \\\n",
" --master-port=\"$COSMOS3_NANO_TEXT_MASTER_PORT\" \\\n",
" -m cosmos_framework.scripts.inference \\\n",
" --parallelism-preset=throughput \\\n",
" --dp-shard-size=\"$FP8_NUM_GPUS\" \\\n",
" --dp-replicate-size=1 \\\n",
" --cp-size=1 \\\n",
" --cfgp-size=1 \\\n",
" --no-use-torch-compile \\\n",
" -i \"$COSMOS3_INPUT_DIR/nano_text.json\" \\\n",
" -o \"$COSMOS3_OUTPUT_ROOT/cosmos_framework_nano_fp8_text\" \\\n",
" --checkpoint-path \"$FP8_CHECKPOINT\" \\\n",
" --seed=0 \\\n",
" --benchmark\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7fdbcd0f",
"metadata": {},
"outputs": [],
"source": [
"from pathlib import Path\n",
"import os, json\n",
"\n",
"output_root = Path(os.environ[\"COSMOS3_OUTPUT_ROOT\"])\n",
"text_path = output_root / \"cosmos_framework_nano_fp8_text\" / \"nano_text\" / \"reasoner_text.txt\"\n",
"benchmark_path = output_root / \"cosmos_framework_nano_fp8_text\" / \"benchmark.json\"\n",
"\n",
"print(text_path)\n",
"print(text_path.read_text())\n",
"if benchmark_path.exists():\n",
" print(json.dumps(json.loads(benchmark_path.read_text()).get(\"average\", {}), indent=2))\n"
]
},
{
"cell_type": "markdown",
"id": "5e7b2393",
Expand Down