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

Đoạn mã FP8 cần thay đổi nằm trong docs/source/en/api/pipelines/cosmos3.md, khoảng dòng 56-69; hướng dẫn ModelOpt mà nó trỏ tới nằm trong docs/source/en/quantization/modelopt.md, khoảng dòng 98-100. Thêm import fp8_per_tensor_gemm và lệnh gọi enable_huggingface_checkpointing() vào đoạn mã, mỗi lệnh kèm một chú thích ngắn. Hoàn thành khi đoạn mã hiển thị cả hai dòng trước from_pretrained và bản build tài liệu hiển thị đúng.

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:

https://github.com/huggingface/diffusers/blob/086bf9578c0f4acbc66e48cd1e7cc26befd9e10f/docs/source/en/api/pipelines/cosmos3.md?plain=1#L56-L69

The ModelOpt guide adds enable_huggingface_checkpointing():

https://github.com/huggingface/diffusers/blob/086bf9578c0f4acbc66e48cd1e7cc26befd9e10f/docs/source/en/quantization/modelopt.md?plain=1#L98-L100

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