Skip to content

cookbook(diffusers): add FP8 quantized checkpoint support - #298

Open
ConstBob wants to merge 5 commits into
NVIDIA:mainfrom
ConstBob:feat/diffusers-fp8
Open

ConstBob wants to merge 5 commits into
NVIDIA:mainfrom
ConstBob:feat/diffusers-fp8

Conversation

@ConstBob

@ConstBob ConstBob commented Oct 21, 2026 •

Copy link
Copy Markdown
Contributor

Do not merge. Wait for legal approval on the FP8 Hugging Face model cards. The model cards need to be approved and published on Hugging Face before this GitHub PR is merged.

Summary

Ports the Diffusers FP8 cookbook support from NVIDIA/cosmos#304 to
NVIDIA/cosmos-framework following the cookbook repository reorganization.
Updates the audiovisual, action forward/inverse dynamics, and transfer notebooks to:

  • select BF16 or FP8 through COSMOS3_DIFFUSERS_PRECISION
  • load Nano, Super, and Edge from their dedicated FP8 repositories
  • load distilled T2I/I2V checkpoints from their fp8 revisions
  • enable ModelOpt checkpoint restoration and FP8 GEMM support
  • explicitly fetch transformer/modelopt_state.pth for a fresh Hub cache
  • preload modular transformers where required by ModelOpt 0.44
  • validate FP8 weights, quantizers, and unresolved meta tensors after loading
    The existing BF16 behavior remains the default.

Repository migration

The original implementation and review history are in NVIDIA/cosmos#304.
After NVIDIA/cosmos#384 and #292 moved cookbook
content into .agents/skills/, the changes were rebased onto the new
Framework locations while preserving the current Framework setup,
uv 0.11.29 guidance, and updated action imports.

Validation

Validated on an NVIDIA RTX PRO 6000 Blackwell:

  • Nano FP8 from an empty Hub cache:
    • ModelOpt state restored
    • 504 FP8 weights, 1533 quantizers, 0 meta tensors
    • complete 189-frame, 1280×720 T2V generation
  • Super FP8 from an independent empty cache:
    • 896 FP8 weights, 2709 quantizers, 0 meta tensors
  • Distilled Super T2I FP8:
    • complete 768×768 image generation
  • Distilled Super I2V FP8:
    • complete 189-frame, 832×480 video generation
  • Nano action forward dynamics, inverse dynamics, and transfer paths were
    exercised end to end in the original implementation.
  • Generated videos were decoded and checked for expected frame count and
    resolution.

Port the Diffusers FP8 loading path onto the cookbooks now under .agents/skills, including ModelOpt state prefetch and quantizer validation.

@pengcuo pengcuo left a comment

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.

Reviewed the four notebooks at cell level. The loading logic looks right to me; everything below is non-blocking except the first point, which is your own hold.

  • Hold stands. nvidia/Cosmos3-{Nano,Super,Edge}-FP8 and the fp8 revisions of the two 4-Step cards resolve with an NVIDIA token but return 401 anonymously as of today, so the notebooks would fail for external users until the cards are public. Please keep the "Do not merge" note until then.
  • I checked the concern that load_components() after update_components(transformer=...) would reload the transformer from repo+subfolder (where ModelOpt ignores the subfolder): diffusers ModularPipeline.load_components(names=None) only loads components whose attribute is still None, so the preloaded FP8 transformer is kept. A one-line comment saying so next to the preload would save the next reader the same check.
  • validate_fp8_transformer counts float8_e4m3fn parameters and modules with "quantizer" in their name. That shows ModelOpt restored its graph, not that the W8A8 GEMM path is active (the fp8_per_tensor_gemm import is what enables it). Fine as a smoke check; just do not describe it as verifying FP8 GEMM support.
  • The same helpers (precision_mode, resolve_pretrained_source, ensure_modelopt_state, validate_fp8_transformer) now live in four notebooks with slightly different signatures (materialize_checkpoint_dir(src, from_kwargs) vs (src, token)). That is the cookbook convention, so no change requested, but it is where drift will come from.

@yzhautouskay

Copy link
Copy Markdown

FYI I have created small PR to diffusers upstream to filter out weight_scale/input_scale warnings of "unused weights" when initializing Cosmos3OmniTransformer. They are supposed to be ignored by Diffusers/Transformers since this is vLLM-Omni/SGLang format, and these warnings pollute the logs

@yzhautouskay

Copy link
Copy Markdown

Hold stands. nvidia/Cosmos3-{Nano,Super,Edge}-FP8 and the fp8 revisions of the two 4-Step cards resolve with an NVIDIA token but return 401 anonymously as of today, so the notebooks would fail for external users until the cards are public. Please keep the "Do not merge" note until then.

Additionally, Edge-FP8 with the access to the private model repo still won't work, because checkpoint wasn't yet published to its HF repo main

Move the audiovisual support check ahead of release_pipe, and describe the FP8 weight check as a restore smoke check rather than GEMM verification.
@ConstBob

Copy link
Copy Markdown
Contributor Author

Thanks for the review. The code comments are addressed in 9956d47.

  • The "Do not merge" note stays until the FP8 Hugging Face model cards are public.
  • The audiovisual FP8 support check now runs inside resolve_pretrained_source, before release_pipe(), matching the action notebooks. An unsupported model raises without dropping a pipeline that is already loaded.
  • Noted next to the distilled and transfer preloads that load_components(names=None) only loads components that are still unset, so the preloaded FP8 transformer is kept.
  • validate_fp8_transformer is documented as a smoke check of FP8 weights, quantizer modules, and meta tensors. It does not verify that the W8A8 GEMM path is active. That path is enabled by importing fp8_per_tensor_gemm.
  • Left the duplicated helpers as they are.
  • Noted the unused-weight warning filter in Fix: filter Cosmos3 warnings for auxiliary FP8 export scales huggingface/diffusers#15107; no change needed in this PR for that.

@ConstBob
ConstBob requested a review from pengcuo October 26, 2026 16:43
Distilled checkpoints share one repository, so a cache keyed only by repo and precision reused the previous pipeline after COSMOS3_FP8_REVISION changed.
@ConstBob

Copy link
Copy Markdown
Contributor Author

The audiovisual get_pipe() cache key now includes COSMOS3_FP8_REVISION before the cache-hit check (94a8450). Distilled FP8 checkpoints share one repository, so a key of only repo and precision returned the pipeline loaded for the previous revision.

Verified with the notebook's get_pipe:

  • The same revision remains a cache hit.
  • Changing the revision releases the cached pipeline and loads the new revision.
  • A nonexistent revision raises instead of returning the previous pipeline.
  • Nano, Super, Edge, and BF16 do not use this revision, and their cache behavior is unchanged.

Loaded nvidia/Cosmos3-Super-Text2Image-4Step at revision fp8 on an RTX PRO 6000: 896 FP8 weights, 2709 quantizers, 0 meta tensors. Repeating the call returned the same pipeline. Switching to a nonexistent revision released that pipeline and failed the load.

@yaoxu-crypto yaoxu-crypto added this to the FP8 Release milestone Oct 30, 2026
@yaoxu-crypto yaoxu-crypto added the Skill New skill or updating existing skill label Oct 30, 2026
@lfengad

lfengad commented Oct 31, 2026

Copy link
Copy Markdown
Collaborator

Thanks for the FP8 cookbook work. This is not ready to merge yet: the existing legal/model-card publication hold remains in place, and GitHub currently reports merge conflicts with main. Could you please retain the hold until the owner confirms approval/publication, then reconcile the branch with the generate-media skill migration without restoring retired notebook paths? Please let us know when both the hold and conflicts are cleared so the final diff can be reviewed.

…kill.

The audiovisual notebook was removed by the skill migration. Keep that path deleted and carry the FP8 loader in the reference that replaced it.
@ConstBob

ConstBob commented Nov 2, 2026

Copy link
Copy Markdown
Contributor Author

Thanks @lfengad .

Merged main at 3b6bf35. The conflict with the generate-media migration is resolved, and the retired audiovisual notebook was not restored.

FP8 loading for that path is now in .agents/skills/cosmos3-generate-media/references/diffusers.md (Quantized Checkpoints). SKILL.md installs nvidia-modelopt==0.44.0 and points to that section. The action and transfer Diffusers notebooks that remain on main were left as they are.

The publication hold is unchanged. This is not ready to merge until the owner confirms the model cards are approved and published.

@satishra-ai satishra-ai left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

PR #298 testing completed: PARTIAL PASS

  • Initial notebook matrix: 14/14 passed.
  • Updated Skill at 24372d4: 5 passed; 2 required a workaround.
  • All 24 outputs validated. Conservative usage: 2.14 / 6 GPU-hours.
  • Filed Nvbug 7034691 for the distilled payload KeyError: num_steps. It is inherited from main, not introduced by the PR’s FP8 changes.

@lfengad

@ConstBob

ConstBob commented Nov 3, 2026

Copy link
Copy Markdown
Contributor Author

Thanks @satishra-ai . The two workaround cases are the distilled create_payload KeyError in the generate-media helper. That helper is unchanged from main, and this PR does not modify it.

The publication hold is unchanged. This is not ready to merge until the owner confirms the model cards are approved and published.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Skill New skill or updating existing skill

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants