Cosmos 3 docs: FP8 snippet fails at the first W8A8 step without importing ModelOpt's FP8 GEMM backend

Describe the bug

The FP8 section of the Cosmos 3 docs loads the ModelOpt FP8 checkpoint with from_pretrained(..., revision="fp8", device_map="cuda") and points to the ModelOpt guide for the restore path:

diffusers/docs/source/en/api/pipelines/cosmos3.md

Lines 56 to 69 in 086bf95

    
           Load the `fp8` revision with the same restore path as the [ModelOpt guide](../../quantization/modelopt) (`revision="fp8"` already carries the quantization config). Mixed precision then follows the checkpoint automatically: 
        
           ```python 
        
           import torch 
        
           from diffusers import Cosmos3OmniPipeline 
        
           pipe = Cosmos3OmniPipeline.from_pretrained( 
        
               "nvidia/Cosmos3-Nano", 
        
               revision="fp8", 
        
               dtype=torch.bfloat16, 
        
               device_map="cuda", 
        
           ) 
        
           result = pipe(prompt="...", num_inference_steps=35) 
        
           ```

The ModelOpt guide adds enable_huggingface_checkpointing():

diffusers/docs/source/en/quantization/modelopt.md

Lines 98 to 100 in 086bf95

    
           from modelopt.torch.opt import enable_huggingface_checkpointing 
        
           enable_huggingface_checkpointing()

With both, a multi-step video generation runs the first W8A16 steps and then fails at the first W8A8 step. ModelOpt finds no real-quant GEMM for the restored RealQuantLinear layers ("RealQuantLinear: No real-quant GEMM found"), builds its fallback extension, and its tensor quantizer raises. The FP8 GEMM backend is registered only by importing modelopt.torch.quantization.backends.fp8_per_tensor_gemm, which neither page shows. NVIDIA's Cosmos cookbook does that import before loading:

https://github.com/NVIDIA/cosmos/blob/fe9d22d601457ff491a3e3265a949083f9c562e3/cookbooks/cosmos3/generator/diffusers_cpu_offloading/cosmos3_streaming/fp8.py#L71-L80

Without enable_huggingface_checkpointing() (the snippet as shown), the FP8 weights load into plain nn.Linear and the first denoising step raises ValueError: Cosmos3 mixed precision found no compatible serialized ModelOpt FP8 linears under ['reasoner', 'generation']; discovered counts={'reasoner': 0, 'generation': 0}.

Unchanged on main: docs/source/en/api/pipelines/cosmos3.md mentions neither call, and no docs page mentions fp8_per_tensor_gemm.

Reproduction

import torch
from diffusers import Cosmos3OmniPipeline
from modelopt.torch.opt import enable_huggingface_checkpointing

enable_huggingface_checkpointing()
pipe = Cosmos3OmniPipeline.from_pretrained("nvidia/Cosmos3-Nano", revision="fp8", dtype=torch.bfloat16, device_map="cuda")
video = pipe(
    prompt="A robot arm stacks colored blocks on a wooden table.",
    num_frames=49, height=480, width=848, num_inference_steps=8,
    generator=torch.Generator(device="cuda").manual_seed(123),
).video

(Our run used the prompt and negative prompt from nvidia/Cosmos3-Nano assets/; the failure happens at step 3 of 8, the first W8A8 step, independently of the prompt.)

Logs

RealQuantLinear: No real-quant GEMM found: RealQuantQuantLinear(
  in_features=4096, out_features=4096, bias=False
  (input_quantizer): TensorQuantizer((4, 3) bit fake per-tensor amax=0.5430 calibrator=MaxCalibrator quant)
  (output_quantizer): TensorQuantizer(disabled)
  (weight_quantizer): TensorQuantizer((4, 3) bit per-tensor amax=0.5039 calibrator=MaxCalibrator quant)
).
Loading extension modelopt_cuda_ext_fp8...
Loaded extension modelopt_cuda_ext_fp8 in 67.1 seconds
Traceback (most recent call last):
  File "repro.py", line 51, in <module>
    video = pipe(
            ^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py", line 1756, in __call__
    preds_vision, preds_sound, preds_action = self.transformer(
                                              ^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/diffusers/models/transformers/transformer_cosmos3.py", line 818, in forward
    und_out, gen_out = self._run_decoder_stack(und_seq, gen_seq, rotary_emb)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/diffusers/models/transformers/transformer_cosmos3.py", line 662, in _run_decoder_stack
    und_seq, gen_seq = decoder_layer(und_seq, gen_seq, rotary_emb)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/diffusers/models/transformers/transformer_cosmos3.py", line 364, in forward
    und_attn_out, gen_attn_out = self.self_attn(und_norm, gen_norm, rotary_emb)
                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/diffusers/models/transformers/transformer_cosmos3.py", line 301, in forward
    return self.processor(self, und_seq, gen_seq, rotary_emb)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/diffusers/models/transformers/transformer_cosmos3.py", line 68, in __call__
    q_und = attn.to_q(und_seq).view(-1, attn.num_attention_heads, attn.head_dim)
            ^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/diffusers/pipelines/cosmos/mixed_precision.py", line 240, in mixed_forward
    return __base(inputs, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/modelopt/torch/quantization/nn/modules/quant_linear.py", line 240, in forward
    return super().forward(input, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/modelopt/torch/quantization/nn/modules/quant_module.py", line 240, in forward
    return super().forward(input, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/modelopt/torch/quantization/nn/modules/quant_module.py", line 192, in forward
    output = super().forward(input, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/linear.py", line 134, in forward
    return F.linear(input, self.weight, self.bias)
                           ^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/modelopt/torch/opt/dynamic.py", line 782, in __getattr__
    return manager.get_da_cb(name)(self, value)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/modelopt/torch/opt/dynamic.py", line 76, in __call__
    val = cb(self_module, val)
          ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/modelopt/torch/quantization/nn/modules/quant_module.py", line 230, in _get_quantized_weight
    return module.weight_quantizer(weight)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/modelopt/torch/quantization/nn/modules/tensor_quantizer.py", line 1105, in forward
    raise ValueError(
ValueError: self._dequantize is True and self.fake_quant is False. This case should have been handled.

With import modelopt.torch.quantization.backends.fp8_per_tensor_gemm added before loading (otherwise the same load), generation completes and the W8A8 steps compile and use ModelOpt's FP8 GEMM. We verified that with enable_group_offload on a 16 GB GPU (35 steps × 189 frames and 8 steps × 49 frames); the traceback above is from the device_map="cuda" run on a 32 GB GPU.

Suggested fix

Add both lines to the Cosmos 3 FP8 snippet:

import modelopt.torch.quantization.backends.fp8_per_tensor_gemm  # noqa: F401 - registers ModelOpt's FP8 GEMM
from modelopt.torch.opt import enable_huggingface_checkpointing

enable_huggingface_checkpointing()  # restores the ModelOpt FP8 linears while loading

Happy to open a docs PR.

System Info

diffusers 0.41.0, nvidia-modelopt 0.44.0, transformers 5.20.2, accelerate 1.15.0
torch 2.9.0a0+145a3a7 (container nvcr.io/nvidia/pytorch:25.10-py3), CUDA 13.0, Python 3.12
NVIDIA GeForce RTX 5090 32 GB, driver 620.62

I hit exactly this on nvidia/Cosmos3-Nano @fp8 and prepared a docs fix — per the contribution guide I'd like an explicit acknowledgment here before opening the PR.

Root cause — each page is one prerequisite short:

cosmos3.md's FP8 snippet calls from_pretrained(revision="fp8") without enable_huggingface_checkpointing(), so the quantized linears restore as plain nn.Linear → ValueError: Cosmos3 mixed precision found no compatible serialized ModelOpt FP8 linears.
modelopt.md's deserialize snippet restores the layout but neither page imports modelopt.torch.quantization.backends.fp8_per_tensor_gemm, so the first W8A8 step has no registered kernel → RealQuantLinear: No real-quant GEMM found.

Fix (2 files, +8/−1): both snippets now do what NVIDIA's Cosmos cookbook does before loading (activate_modelopt_fp8() at fe9d22d does the same two calls) — call enable_huggingface_checkpointing() before from_pretrained and import the backend module once. The cosmos3 prose now names both failure modes so these errors are searchable from the docs.

Branch: https://github.com/Shanmukh-l33t/diffusers/commits/docs-fp8-modelopt-restore-15170 (commit cb042fba3b5f73371b53357a7bb0c43fa6339a4b)
Compare: main...Shanmukh-l33t:docs-fp8-modelopt-restore-15170

Verification: a docs-consistency checker over both files fails on main with 3 findings (the two missing cosmos3 prerequisites + the missing backend import in modelopt.md) and passes after the edit; every python block in both files still parses with ast.parse. Scope check: revision="fp8" / quant_type FP8 loading appears only in these two docs pages, so no other snippet needs the same treatment.

If you'd rather this land differently (only cosmos3.md, a shared snippet, or a different wording), say so and I'll adjust before opening the PR.

Navigation Menu

Repository navigation

Cosmos 3 docs: FP8 snippet fails at the first W8A8 step without importing ModelOpt's FP8 GEMM backend #15170

Describe the bug

Reproduction

Logs

Suggested fix

System Info

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	Load the `fp8` revision with the same restore path as the [ModelOpt guide](../../quantization/modelopt) (`revision="fp8"` already carries the quantization config). Mixed precision then follows the checkpoint automatically:

	```python
	import torch
	from diffusers import Cosmos3OmniPipeline

	pipe = Cosmos3OmniPipeline.from_pretrained(
	"nvidia/Cosmos3-Nano",
	revision="fp8",
	dtype=torch.bfloat16,
	device_map="cuda",
	)
	result = pipe(prompt="...", num_inference_steps=35)
	```

	from modelopt.torch.opt import enable_huggingface_checkpointing

	enable_huggingface_checkpointing()

Cosmos 3 docs: FP8 snippet fails at the first W8A8 step without importing ModelOpt's FP8 GEMM backend #15170

Description

Describe the bug

Reproduction

Logs

Suggested fix

System Info

Activity

Shanmukh-l33t commented on Nov 7, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions