Cosmos 3 FP8: wrapping pipe.transformer (e.g. torch.compile) silently disables the checkpoint's W8A16 step policy

Describe the bug

For the ModelOpt FP8 Cosmos 3 checkpoints, the per-step W8A16/W8A8 policy (quantization_config.runtime.diffusion_step_policy in transformer/config.json, added in #14664) is resolved from self.transformer on every pipeline call:

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

Lines 1709 to 1715 in 086bf95

    
           mixed_precision = Cosmos3MixedPrecisionConfig.resolve( 
        
               self.transformer, 
        
               mixed_precision_format=mixed_precision_format, 
        
               mixed_precision_first_steps=mixed_precision_first_steps, 
        
               mixed_precision_last_steps=mixed_precision_last_steps, 
        
               mixed_precision_reasoner_policy=mixed_precision_reasoner_policy, 
        
           )

Cosmos3MixedPrecisionConfig.resolve() → quantization_config_from_module() → _on_disk_quantization_config(), which looks up the config loader on the module's class:

diffusers/src/diffusers/pipelines/cosmos/mixed_precision.py

Lines 443 to 464 in 086bf95

    
           def _on_disk_quantization_config(module: nn.Module) -> dict[str, Any] | None: 
        
               config = getattr(module, "config", None) 
        
               name_or_path = None 
        
               if config is not None: 
        
                   name_or_path = getattr(config, "_name_or_path", None) 
        
                   if name_or_path is None and hasattr(config, "get"): 
        
                       name_or_path = config.get("_name_or_path") 
        
               loader = getattr(type(module), "load_config", None) 
        
               if not name_or_path or not callable(loader): 
        
                   return None 
        
               try: 
        
                   raw = loader(name_or_path, local_files_only=True) 
        
               except TypeError: 
        
                   try: 
        
                       raw = loader(name_or_path) 
        
                   except Exception: 
        
                       raw = None 
        
               except Exception: 
        
                   raw = None 
        
               if isinstance(raw, dict): 
        
                   return _maybe_mapping(raw.get("quantization_config")) 
        
               return None

loader = getattr(type(module), "load_config", None)

The in-memory config does not carry quantization_config.runtime (see the docstring of quantization_config_from_module, L409-L426), so this on-disk read is the only source of the policy. When pipe.transformer is a wrapper, type(module) is the wrapper class, which has no load_config. The on-disk config is never read, _parsed_checkpoint_policy() returns None, and resolve() returns format="none" (L157-L158). Every step then runs W8A8 instead of the checkpoint's W8A16 first/last steps, with no warning. The Cosmos 3 docs note that all-W8A8 multi-step video can flicker (cosmos3.md L50-L54).

pipe.transformer = torch.compile(pipe.transformer), the pattern shown in several docs pages (e.g. optimization/cache_dit.md L260, quantization/autoround.md L102), is such a wrapper (OptimizedModule). In-place pipe.transformer.compile() and compile_repeated_blocks() (used on the Cosmos 3 page) keep the class and are not affected. The line is unchanged on main: mixed_precision.py#L450.

Reproduction

The policy is printed with the same resolve() call the pipeline makes, before any generation:

import modelopt.torch.quantization.backends.fp8_per_tensor_gemm  # noqa: F401 - registers the FP8 GEMM
import torch
from diffusers import Cosmos3OmniPipeline
from diffusers.pipelines.cosmos.mixed_precision import Cosmos3MixedPrecisionConfig
from modelopt.torch.opt import enable_huggingface_checkpointing

enable_huggingface_checkpointing()
pipe = Cosmos3OmniPipeline.from_pretrained("nvidia/Cosmos3-Nano", revision="fp8", dtype=torch.bfloat16)
pipe.vae.to("cuda")
pipe.transformer.enable_group_offload(
    onload_device=torch.device("cuda"), offload_device=torch.device("cpu"),
    offload_type="block_level", num_blocks_per_group=1, use_stream=True,
)

pipe.transformer = torch.compile(pipe.transformer, backend="eager")  # the only difference between the two runs
print(f"type(pipe.transformer) = {type(pipe.transformer).__name__}")
print(f"resolved policy = {Cosmos3MixedPrecisionConfig.resolve(pipe.transformer)}")

Logs

Without the torch.compile line:

type(pipe.transformer) = Cosmos3OmniTransformer
resolved policy = Cosmos3MixedPrecisionConfig(format='fp8', first_steps=3, last_steps=3, reasoner_policy='high_precision', overlap='a16')

With it:

type(pipe.transformer) = OptimizedModule
resolved policy = Cosmos3MixedPrecisionConfig(format='none', first_steps=3, last_steps=3, reasoner_policy='high_precision', overlap='a16')

Effect on generation (T2V 848×480, 49 frames, 8 steps, seed 123, transformer with enable_group_offload(offload_type="block_level", num_blocks_per_group=1, use_stream=True) to fit a 16 GB GPU):

`pipe.transformer`	call	resolved policy	mixed-precision runtime installed on the transformer
`Cosmos3OmniTransformer`	default	`fp8`, first/last 3 steps W8A16	yes (252 reasoner + 252 generation linears)
`Cosmos3OmniTransformer`	`mixed_precision_format="none"`	`none`	no; different video from the default run
`OptimizedModule` (`torch.compile(pipe.transformer, backend="eager")`)	default	`none`	— (see below)

With the compiled ModelOpt FP8 transformer, generation currently stops earlier inside Dynamo, so the end-to-end output of the compiled run cannot be shown; the resolved policy above is what the pipeline applies to it. Any wrapper that forwards attribute access and does run reaches the none schedule.

Dynamo traceback of the compiled run (separate issue, context only)

Traceback (most recent call last):
  File "repro.py", line 73, 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/_dynamo/eval_frame.py", line 414, in __call__
    return super().__call__(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  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/torch/_dynamo/eval_frame.py", line 832, in compile_wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  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/hooks/hooks.py", line 217, in new_forward
    output = function_reference.forward(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/diffusers/hooks/hooks.py", line 216, in new_forward
    args, kwargs = function_reference.pre_forward(module, *args, **kwargs)
  [... torch/_dynamo frames ...]
  File "/usr/local/lib/python3.12/dist-packages/modelopt/torch/opt/dynamic.py", line 81, in __len__
    return len(self._callbacks)
               ^^^^^^^^^^^^^^^
torch._dynamo.exc.InternalTorchDynamoError: AttributeError: '_FoldedCallback' object has no attribute '_callbacks'

from user code:
   File "/usr/local/lib/python3.12/dist-packages/diffusers/hooks/hooks.py", line 217, in torch_dynamo_resume_in_new_forward_at_216
    output = function_reference.forward(*args, **kwargs)
  File "/usr/local/lib/python3.12/dist-packages/diffusers/models/transformers/transformer_cosmos3.py", line 745, in forward
    packed_tokens_vision = self.proj_in(packed_tokens_vision)
  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)

Suggested fix

Either of these, plus a regression test (resolve(torch.compile(model)) == resolve(model) on a small model whose config declares the ModelOpt FP8 policy):

Look the loader up on the instance: loader = getattr(module, "load_config", None). load_config is a ConfigMixin classmethod, so this returns the same bound classmethod, and wrappers that forward attribute access (OptimizedModule.__getattr__ → _orig_mod) work too.
Or unwrap first with the existing helper unwrap_module, as other Diffusers code does for compiled modules. This covers torch.compile only.

Optionally, a warning when a ModelOpt FP8 backend is detected but the on-disk policy could not be read, so the fallback to all-W8A8 is visible. Happy to open a 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 5080 16 GB, driver 620.62

Navigation Menu

Repository navigation

Cosmos 3 FP8: wrapping pipe.transformer (e.g. torch.compile) silently disables the checkpoint's W8A16 step policy #15169

Describe the bug

Reproduction

Logs

Suggested fix

System Info

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	mixed_precision = Cosmos3MixedPrecisionConfig.resolve(
	self.transformer,
	mixed_precision_format=mixed_precision_format,
	mixed_precision_first_steps=mixed_precision_first_steps,
	mixed_precision_last_steps=mixed_precision_last_steps,
	mixed_precision_reasoner_policy=mixed_precision_reasoner_policy,
	)

	def _on_disk_quantization_config(module: nn.Module) -> dict[str, Any] \| None:
	config = getattr(module, "config", None)
	name_or_path = None
	if config is not None:
	name_or_path = getattr(config, "_name_or_path", None)
	if name_or_path is None and hasattr(config, "get"):
	name_or_path = config.get("_name_or_path")
	loader = getattr(type(module), "load_config", None)
	if not name_or_path or not callable(loader):
	return None
	try:
	raw = loader(name_or_path, local_files_only=True)
	except TypeError:
	try:
	raw = loader(name_or_path)
	except Exception:
	raw = None
	except Exception:
	raw = None
	if isinstance(raw, dict):
	return _maybe_mapping(raw.get("quantization_config"))
	return None

Cosmos 3 FP8: wrapping pipe.transformer (e.g. torch.compile) silently disables the checkpoint's W8A16 step policy #15169

Description

Describe the bug

Reproduction

Logs

Suggested fix

System Info

Activity

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions