iframe-proxy

ConstBob · 2026-08-31T21:07:12Z

What does this PR do?

Adds denoising-step-aware mixed W8A8/W8A16 for Cosmos3 ModelOpt FP8 checkpoints, so official Hub revision=fp8 is not W8A8 on every step.

This is the Diffusers counterpart of the FP8 schedule contract in vllm-omni#6560 (not a port of NVFP4, tensor parallel, caches, or vllm serve).

Behavior

Checkpoint weights stay serialized static W8A8 (weight_only: false). Mixing is runtime.
Video Nano / Super / Super-I2V fp8: first 3 + last 3 scheduler steps W8A16 (dequant + F.linear); middle steps native ModelOpt W8A8.
Precision is chosen once per scheduler step (CFG cond/uncond match). Reasoner defaults to A16 when the checkpoint says reasoner: a16.
Distilled 4-step and Super-T2I fp8 have runtime: null → mixed off, all steps native W8A8.
mixed_precision_format=None (default) reads quantization_config.runtime.diffusion_step_policy from the transformer (including on-disk transformer/config.json when the live ModelOpt config omits runtime). "none" disables only the schedule. "fp8" forces it.

Code

src/diffusers/pipelines/cosmos/mixed_precision.py — policy parse (fail-closed), linear discovery, W8A16 dispatch.
pipeline_cosmos3_omni.py and modular denoise.py — apply/reset once per step.
docs/source/en/api/pipelines/cosmos3.md — user-facing FP8 mixed section.
Out of scope: NVFP4, custom kernels, dense-weight cache, cookbook/Hub cards (follow-up after this lands).

Tests

Unit (CPU, in-tree): tests/pipelines/cosmos/test_cosmos3_mixed_precision.py — 17 passed (schedule boundaries, 1-step stays native, reasoner vs generation, overlay of live runtime: null + on-disk policy, FP32 W8A16, malformed policy / NVFP4 reject).

pytest tests/pipelines/cosmos/test_cosmos3_mixed_precision.py
# 17 passed, 5 subtests passed

GPU, not in CI (official Hub revision=fp8):

Cosmos3-Nano @fp8: auto-enable 3+3; 8-step T2I E2E with exact trace W8A16×3 / W8A8×2 / W8A16×3; 35-step latent W8A16×3 / W8A8×29 / W8A16×3.
Super-Text2Image-4Step @fp8 transformer: auto-off; GEMM stays native W8A8.

Not claimed in this PR: Super 64B generate, 189-frame I2V, vs-BF16 hashes, or the vLLM-Omni H100 throughput table.

Depends on NVIDIA ModelOpt restore for serialized FP8 (same as existing Cosmos3 FP8 loading).

Self-review (final)

Ran the self-review skill on huggingface/main...HEAD against .ai/references/review-rules.md (plus pipelines.md, modular.md, testing.md, code_style.md, pitfalls.md).

Diff files: mixed_precision.py, pipeline_cosmos3_omni.py, modular_pipelines/cosmos/denoise.py, test_cosmos3_mixed_precision.py, docs/source/en/api/pipelines/cosmos3.md.

Blocking issues: none.

Non-blocking — left for review (deliberate):

Transformer forward wrappers stay installed after __call__ (mixed_precision.py). Inactive wrappers call the original ModelOpt forward. Per pipelines.md gotcha 7 we could uninstall on exit (PAG-style); we kept wrappers to avoid re-wrapping hundreds of linears each generate.
W8A16 reads layer.weight then F.linear. Per testing.md, that bypasses group-offload leaf hooks. Same dense A16 reference as vLLM-Omni; group offload + serialized FP8 was not a target.
Focused tests use fake ModelOpt linears + unittest.TestCase. Per testing.md, new pipeline tests prefer real tiny classes and pytest mixins. Existing Cosmos3 pipeline tests are still unittest; ModelOpt restore is impractical in CI.
mixed_precision_overlap is parsed from the checkpoint / resolve(), not exposed on Cosmos3OmniPipeline.__call__.
Docs revision="fp8" example does not repeat full ModelOpt from_pretrained kwargs; loading still follows the ModelOpt guide.

Dead code: FIRST_LAST_N_FP8_POLICY, from_kwargs, and quantization_config_from_module are used (tests + resolve). No unused runtime methods.

Verdict: READY.

Before submitting

Who can review?

@yiyixuxu @asomoza

…oints. Keep native ModelOpt GEMM on middle steps and dequant-linear W8A16 on the first/last steps so CFG cond/uncond share one precision per scheduler step.

Enable first/last W8A16 only when transformer/config.json declares diffusion_step_policy, so distilled FP8 stays native W8A8 instead of inheriting a hardcoded 3+3 window.

…schedules. Read the checkpoint runtime policy from on-disk transformer/config.json when the live ModelOpt config omits it, fail closed on incomplete policies, and allow FP32 activations on the W8A16 path.

sayakpaul

Thanks for the PR. I think it's a bit intrusive in terms of what it is doing. Could you provide runtime numbers (latency, throughput) and also how does it compare in terms of visual results (with and without comparisons would be great).

… drop unit tests. W8A8/W8A16 is explained without listing call-site overrides; TorchAO and other quantizers keep their native forwards; focused tests are removed until Hub fp8 usage is clearer.

ConstBob · 2026-09-08T19:57:46Z

sayakpaul

Thanks for the iteration. Do you think it could make sense to try to think of an approach where we make the changes bit more general, i.e., we make it so that these are done at the Modelopt quantizer class?

diffusers/src/diffusers/quantizers/modelopt/modelopt_quantizer.py

Line 28 in a71e62e

class NVIDIAModelOptQuantizer(DiffusersQuantizer):

Or do you think the pipeline-level call arguments make it difficult to do so?

…e boilerplate. Keep serialized restore in the ModelOpt guide and spell out that all-W8A8 is faster but can flicker on multi-step video.

ConstBob · 2026-09-11T18:20:22Z

Thanks @sayakpaul for the review. I've updated the PR.

I looked at the other Cosmos3 FP8 runtimes that already landed this schedule. None of them put it on the generic ModelOpt loader/quantizer.

vLLM-Omni (vllm-project/vllm-omni#6560, merged) keeps policy parse + W8A16 dispatch under diffusion/models/cosmos3/mixed_precision/. The pipeline selects precision at each denoising-loop iteration; generate-time disable is --additional-config '{"cosmos3_mixed_precision":{"enabled":false}}'.

SGLang (sgl-project/sglang#36380, merged) does the same from the Cosmos3 denoise stage (StepMixedPrecisionController + dequant/F.linear). The generic FP8 linear method is unchanged; an env var disables only the schedule.

That matches this PR: NVIDIAModelOptQuantizer is load-time (restore FP8 weights), while the schedule is a function of scheduler step index, CFG, and reasoner vs generation. The __call__ kwargs are the Diffusers equivalent of vLLM's additional-config / SGLang's env — request-scoped, not from_pretrained.

To align with this PR's counterparts in other runtimes, I'd prefer to keep that split in this PR rather than moving the schedule onto NVIDIAModelOptQuantizer. Does that sound reasonable to you?

sayakpaul · 2026-09-13T14:53:07Z

@bot /style

github-actions · 2026-09-13T14:54:30Z

Style bot fixed some files and pushed the changes.

HuggingFaceDocBuilderDev · 2026-09-13T15:01:30Z

ConstBob and others added 4 commits August 24, 2026 16:42

Add opt-in mixed W8A8/W8A16 denoising for Cosmos3 ModelOpt FP8 checkp…

b94bf63

…oints. Keep native ModelOpt GEMM on middle steps and dequant-linear W8A16 on the first/last steps so CFG cond/uncond share one precision per scheduler step.

Merge branch 'huggingface:main' into feat/cosmos3-mixed-w8a8-w8a16

261265f

Read Cosmos3 mixed-precision schedule from the checkpoint policy.

37c9248

Enable first/last W8A16 only when transformer/config.json declares diffusion_step_policy, so distilled FP8 stays native W8A8 instead of inheriting a hardcoded 3+3 window.

Harden Cosmos3 mixed-precision loading and document official Hub fp8 …

6939c05

…schedules. Read the checkpoint runtime policy from on-disk transformer/config.json when the live ModelOpt config omits it, fail closed on incomplete policies, and allow FP32 activations on the W8A16 path.

github-actions Bot added documentation Improvements or additions to documentation tests modular-pipelines pipelines size/L PR with diff > 200 LOC labels Aug 31, 2026

yiyixuxu requested a review from sayakpaul September 2, 2026 01:00

sayakpaul reviewed Sep 3, 2026

View reviewed changes

Comment thread docs/source/en/api/pipelines/cosmos3.md Outdated

Comment thread docs/source/en/api/pipelines/cosmos3.md Outdated

Comment thread tests/pipelines/cosmos/test_cosmos3_mixed_precision.py Outdated

Comment thread src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py

ConstBob and others added 2 commits September 8, 2026 09:56

Merge branch 'huggingface:main' into feat/cosmos3-mixed-w8a8-w8a16

30c0852

Address review: simplify FP8 mixed docs, no-op non-ModelOpt backends,…

1b0696b

… drop unit tests. W8A8/W8A16 is explained without listing call-site overrides; TorchAO and other quantizers keep their native forwards; focused tests are removed until Hub fp8 usage is clearer.

github-actions Bot removed the tests label Sep 8, 2026

ConstBob requested a review from sayakpaul September 8, 2026 20:13

sayakpaul reviewed Sep 11, 2026

View reviewed changes

Comment thread docs/source/en/api/pipelines/cosmos3.md Outdated

Comment thread docs/source/en/api/pipelines/cosmos3.md Outdated

Document FP8 mixed default vs none trade-offs without ModelOpt restor…

dbaca82

…e boilerplate. Keep serialized restore in the ModelOpt guide and spell out that all-W8A8 is faster but can flicker on multi-step video.

ConstBob requested a review from sayakpaul September 11, 2026 18:22

sayakpaul approved these changes Sep 13, 2026

View reviewed changes

Merge branch 'main' into feat/cosmos3-mixed-w8a8-w8a16

6cd0e1a

Apply style fixes

39ee546

Regenerate Cosmos modular auto docstrings for mixed-precision inputs.

fc4672d

ConstBob requested a review from sayakpaul September 14, 2026 16:23

Merge branch 'main' into feat/cosmos3-mixed-w8a8-w8a16

4cc1bff

sayakpaul merged commit 759164b into huggingface:main Sep 14, 2026
19 checks passed

Sunbelt Computer Software

PL/B Language Development and Support

Navigation Menu

Repository navigation

[Cosmos3] Mixed W8A8/W8A16 denoising for ModelOpt FP8 checkpoints - #14664

[Cosmos3] Mixed W8A8/W8A16 denoising for ModelOpt FP8 checkpoints#14664
sayakpaul merged 11 commits into
huggingface:mainfrom
ConstBob:feat/cosmos3-mixed-w8a8-w8a16

ConstBob commented Aug 31, 2026 •

edited

Loading

Uh oh!

sayakpaul left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ConstBob commented Sep 8, 2026

Uh oh!

sayakpaul left a comment

Uh oh!

Uh oh!

Uh oh!

ConstBob commented Sep 11, 2026

Uh oh!

sayakpaul commented Sep 13, 2026

Uh oh!

github-actions Bot commented Sep 13, 2026 •

edited

Loading

Uh oh!

HuggingFaceDocBuilderDev commented Sep 13, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Sunbelt Computer Software

PL/B Language Development and Support

Conversation

ConstBob commented Aug 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Behavior

Code

Tests

Self-review (final)

Before submitting

Who can review?

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ConstBob commented Sep 8, 2026

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ConstBob commented Sep 11, 2026

Uh oh!

sayakpaul commented Sep 13, 2026

Uh oh!

github-actions Bot commented Sep 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Sep 13, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ConstBob commented Aug 31, 2026 •

edited

Loading

github-actions Bot commented Sep 13, 2026 •

edited

Loading