From b94bf63a3fbc1600ea62efe029ef6d2d2f0f3a58 Mon Sep 17 00:00:00 2001 From: Yiming Zhao Date: Mon, 24 Aug 2026 16:42:20 -0700 Subject: [PATCH 1/3] Add opt-in mixed W8A8/W8A16 denoising for Cosmos3 ModelOpt FP8 checkpoints. 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. --- .../modular_pipelines/cosmos/denoise.py | 64 +++- .../pipelines/cosmos/mixed_precision.py | 255 ++++++++++++++++ .../pipelines/cosmos/pipeline_cosmos3_omni.py | 288 ++++++++++-------- .../cosmos/test_cosmos3_mixed_precision.py | 160 ++++++++++ 4 files changed, 638 insertions(+), 129 deletions(-) create mode 100644 src/diffusers/pipelines/cosmos/mixed_precision.py create mode 100644 tests/pipelines/cosmos/test_cosmos3_mixed_precision.py diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py index eda37c8e99cf..c7003e593dbb 100644 --- a/src/diffusers/modular_pipelines/cosmos/denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/denoise.py @@ -3,6 +3,11 @@ import torch from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer +from ...pipelines.cosmos.mixed_precision import ( + Cosmos3MixedPrecisionConfig, + apply_cosmos3_mixed_precision_step, + reset_cosmos3_mixed_precision, +) from ...schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler from ..modular_pipeline import ( BlockState, @@ -463,18 +468,63 @@ def loop_inputs(self) -> list[InputParam]: InputParam( name="num_warmup_steps", type_hint=int, required=True, description="Number of scheduler warmup steps." ), + InputParam( + name="mixed_precision_format", + type_hint=str, + default="none", + description="Set to 'fp8' to enable mixed W8A8/W8A16 denoising.", + ), + InputParam( + name="mixed_precision_first_steps", + type_hint=int, + default=3, + description="Leading W8A16 step count when mixed precision is enabled.", + ), + InputParam( + name="mixed_precision_last_steps", + type_hint=int, + default=3, + description="Trailing W8A16 step count when mixed precision is enabled.", + ), + InputParam( + name="mixed_precision_reasoner_policy", + type_hint=str, + default="high_precision", + description="Use W8A16 or checkpoint-native W8A8 for the reasoner path.", + ), ] @torch.no_grad() def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: - for i, t in enumerate(block_state.timesteps): - components, block_state = self.loop_step(components, block_state, i=i, t=t) - if i == len(block_state.timesteps) - 1 or ( - (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0 - ): - progress_bar.update() + first_steps = getattr(block_state, "mixed_precision_first_steps", None) + last_steps = getattr(block_state, "mixed_precision_last_steps", None) + mixed_precision = Cosmos3MixedPrecisionConfig.from_kwargs( + mixed_precision_format=getattr(block_state, "mixed_precision_format", None) or "none", + mixed_precision_first_steps=3 if first_steps is None else first_steps, + mixed_precision_last_steps=3 if last_steps is None else last_steps, + mixed_precision_reasoner_policy=getattr(block_state, "mixed_precision_reasoner_policy", None) + or "high_precision", + ) + trace = [] + try: + with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: + for i, t in enumerate(block_state.timesteps): + apply_cosmos3_mixed_precision_step( + components.transformer, + mixed_precision, + i, + len(block_state.timesteps), + trace=trace, + ) + components, block_state = self.loop_step(components, block_state, i=i, t=t) + if i == len(block_state.timesteps) - 1 or ( + (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0 + ): + progress_bar.update() + finally: + reset_cosmos3_mixed_precision(components.transformer, mixed_precision) + components._mixed_precision_trace = trace self.set_block_state(state, block_state) return components, state diff --git a/src/diffusers/pipelines/cosmos/mixed_precision.py b/src/diffusers/pipelines/cosmos/mixed_precision.py new file mode 100644 index 000000000000..7e641d9010da --- /dev/null +++ b/src/diffusers/pipelines/cosmos/mixed_precision.py @@ -0,0 +1,255 @@ +# Copyright 2026 The NVIDIA Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Mixed W8A8/W8A16 denoising for serialized Cosmos3 ModelOpt FP8 checkpoints. + +The checkpoint's restored ModelOpt forward is the native W8A8 path. W8A16 +bypasses that forward, dequantizes the FP8 weight to the activation dtype, and +uses :func:`torch.nn.functional.linear`. This matches the uncached strategy in +vLLM-Omni (vllm-project/vllm-omni#6560). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import MethodType +from typing import Any, Literal + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +MixedPrecisionFormat = Literal["none", "fp8"] +ReasonerPolicy = Literal["high_precision", "base_precision"] +MIXED_PRECISION_FORMATS = frozenset({"none", "fp8"}) +REASONER_POLICIES = frozenset({"high_precision", "base_precision"}) +_RUNTIME_ATTRIBUTE = "_cosmos3_mixed_precision_runtime" + + +@dataclass(frozen=True) +class Cosmos3MixedPrecisionConfig: + """Validated first/last-step W8A16 policy for Cosmos3 FP8 denoising.""" + + format: MixedPrecisionFormat = "none" + first_steps: int = 3 + last_steps: int = 3 + reasoner_policy: ReasonerPolicy = "high_precision" + + @classmethod + def from_kwargs( + cls, + mixed_precision_format: str = "none", + mixed_precision_first_steps: int = 3, + mixed_precision_last_steps: int = 3, + mixed_precision_reasoner_policy: str = "high_precision", + ) -> Cosmos3MixedPrecisionConfig: + precision_format = str(mixed_precision_format).strip().lower() + if precision_format not in MIXED_PRECISION_FORMATS: + raise ValueError( + "mixed_precision_format must be one of " + f"{sorted(MIXED_PRECISION_FORMATS)}, got {mixed_precision_format!r}" + ) + reasoner_policy = str(mixed_precision_reasoner_policy).strip().lower() + if reasoner_policy not in REASONER_POLICIES: + raise ValueError( + "mixed_precision_reasoner_policy must be one of " + f"{sorted(REASONER_POLICIES)}, got {mixed_precision_reasoner_policy!r}" + ) + return cls( + format=precision_format, # type: ignore[arg-type] + first_steps=_non_negative_int(mixed_precision_first_steps, "mixed_precision_first_steps"), + last_steps=_non_negative_int(mixed_precision_last_steps, "mixed_precision_last_steps"), + reasoner_policy=reasoner_policy, # type: ignore[arg-type] + ) + + @property + def enabled(self) -> bool: + return self.format != "none" + + def use_high_precision(self, step_index: int, num_steps: int) -> bool: + """Return True when this scheduler step should run W8A16 (no activation quant).""" + if num_steps <= 0: + raise ValueError(f"num_steps must be positive, got {num_steps}") + if step_index < 0 or step_index >= num_steps: + raise IndexError(f"step_index must be in [0, {num_steps}), got {step_index}") + # Match vLLM-Omni: a 1-step request keeps the checkpoint's base (W8A8) path. + if num_steps == 1: + return False + return step_index < self.first_steps or step_index >= num_steps - self.last_steps + + def precision_name(self, step_index: int, num_steps: int) -> str: + if not self.enabled: + return "base" + return "W8A16" if self.use_high_precision(step_index, num_steps) else "W8A8" + + +class Cosmos3MixedPrecisionRuntime: + """Own the per-step selection and wrapped ModelOpt linear inventory.""" + + def __init__(self, transformer: nn.Module, config: Cosmos3MixedPrecisionConfig) -> None: + self.transformer = transformer + self.config = config + self.active = False + self.generation_high_precision = False + self.installed_counts = {"reasoner": 0, "generation": 0} + self._install() + + def _install(self) -> None: + inventory = [] + for name, layer in self.transformer.named_modules(): + path = _classify_cosmos3_linear(name) + if path is None or not _is_modelopt_fp8_linear(layer): + continue + _validate_modelopt_fp8_linear(name, layer) + inventory.append((name, layer, path)) + self.installed_counts[path] += 1 + + missing = [path for path, count in self.installed_counts.items() if count == 0] + if missing: + raise ValueError( + "Cosmos3 mixed precision found no compatible serialized ModelOpt FP8 linears under " + f"{missing}; discovered counts={self.installed_counts}" + ) + + for name, layer, path in inventory: + original_forward = layer.forward + + def mixed_forward(layer_self, inputs, *args, __name=name, __path=path, __base=original_forward, **kwargs): + if not self.active or not self.use_high_precision(__path): + return __base(inputs, *args, **kwargs) + if args or kwargs: + raise TypeError(f"{__name} W8A16 only supports the standard Linear forward(input) signature") + return _w8a16_linear(layer_self, inputs, __name) + + layer.forward = MethodType(mixed_forward, layer) + + def use_high_precision(self, path: str) -> bool: + if path == "reasoner": + return self.config.reasoner_policy == "high_precision" + return self.generation_high_precision + + def set_step(self, step_index: int, num_steps: int) -> str: + self.active = True + self.generation_high_precision = self.config.use_high_precision(step_index, num_steps) + return "W8A16" if self.generation_high_precision else "W8A8" + + def reset(self) -> None: + self.active = False + self.generation_high_precision = False + + +def _classify_cosmos3_linear(name: str) -> str | None: + """Classify the two MoT paths in Diffusers' fused Cosmos3 decoder layer.""" + if ".mlp_moe_gen." in name or any( + name.endswith(suffix) + for suffix in ( + ".self_attn.to_q", + ".self_attn.to_k", + ".self_attn.to_v", + ".self_attn.to_out", + ) + ): + return "generation" + if ".mlp." in name or any( + name.endswith(suffix) + for suffix in ( + ".self_attn.add_q_proj", + ".self_attn.add_k_proj", + ".self_attn.add_v_proj", + ".self_attn.to_add_out", + ) + ): + return "reasoner" + return None + + +def _is_modelopt_fp8_linear(layer: nn.Module) -> bool: + weight = getattr(layer, "weight", None) + weight_quantizer = getattr(layer, "weight_quantizer", None) + return ( + isinstance(weight, torch.Tensor) + and weight.dtype == torch.float8_e4m3fn + and weight_quantizer is not None + and hasattr(layer, "_should_run_real_quant_gemm") + ) + + +def _validate_modelopt_fp8_linear(name: str, layer: nn.Module) -> None: + weight = layer.weight + input_quantizer = getattr(layer, "input_quantizer", None) + if input_quantizer is not None and not input_quantizer.is_enabled: + raise ValueError(f"{name} has a disabled input quantizer and therefore is not a native W8A8 linear") + if not layer.weight_quantizer.is_enabled: + raise ValueError(f"{name} has a disabled weight quantizer and therefore is not a native W8A8 linear") + scale = getattr(layer.weight_quantizer, "_scale", None) + if not isinstance(scale, torch.Tensor) or scale.numel() != 1: + raise ValueError(f"{name} requires one tensorwise ModelOpt FP8 weight scale") + if weight.ndim != 2: + raise ValueError(f"{name} expected a 2D FP8 weight, got shape {tuple(weight.shape)}") + if hasattr(layer, "pre_quant_scale"): + raise ValueError(f"{name} uses SmoothQuant pre_quant_scale, which mixed W8A8/W8A16 does not support") + + +def _w8a16_linear(layer: nn.Module, inputs: torch.Tensor, name: str) -> torch.Tensor: + if inputs.dtype not in (torch.bfloat16, torch.float16): + raise TypeError(f"{name} W8A16 requires BF16/FP16 activations, got {inputs.dtype}") + scale = layer.weight_quantizer._scale.to(device=layer.weight.device, dtype=inputs.dtype) + dense_weight = layer.weight.to(dtype=inputs.dtype) * scale + return F.linear(inputs, dense_weight, layer.bias) + + +def _get_or_create_runtime(module: nn.Module, config: Cosmos3MixedPrecisionConfig) -> Cosmos3MixedPrecisionRuntime: + runtime = getattr(module, _RUNTIME_ATTRIBUTE, None) + if runtime is None: + runtime = Cosmos3MixedPrecisionRuntime(module, config) + setattr(module, _RUNTIME_ATTRIBUTE, runtime) + elif runtime.config != config: + runtime.config = config + return runtime + + +def apply_cosmos3_mixed_precision_step( + module: nn.Module, + config: Cosmos3MixedPrecisionConfig, + step_index: int, + num_steps: int, + trace: list[str] | None = None, +) -> str: + """Select W8A8 vs W8A16 for one scheduler step. No-op when mixed precision is disabled.""" + name = config.precision_name(step_index, num_steps) + if trace is not None: + trace.append(name) + if not config.enabled: + return name + runtime = _get_or_create_runtime(module, config) + selected = runtime.set_step(step_index, num_steps) + if selected != name: + raise RuntimeError(f"Mixed-precision schedule mismatch: config={name}, runtime={selected}") + return selected + + +def reset_cosmos3_mixed_precision(module: nn.Module, config: Cosmos3MixedPrecisionConfig) -> None: + """Return installed wrappers to the checkpoint's native W8A8 path.""" + if not config.enabled: + return + runtime = getattr(module, _RUNTIME_ATTRIBUTE, None) + if runtime is not None: + runtime.reset() + + +def _non_negative_int(value: Any, name: str) -> int: + if not isinstance(value, int) or isinstance(value, bool) or value < 0: + raise TypeError(f"{name} must be a non-negative integer, got {value!r}") + return value diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py index 589e0ed3d6b0..09309105f60c 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py @@ -36,6 +36,11 @@ from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline +from .mixed_precision import ( + Cosmos3MixedPrecisionConfig, + apply_cosmos3_mixed_precision_step, + reset_cosmos3_mixed_precision, +) logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -1340,6 +1345,10 @@ def __call__( add_resolution_template: bool = True, add_duration_template: bool = True, enable_safety_check: bool = True, + mixed_precision_format: str = "none", + mixed_precision_first_steps: int = 3, + mixed_precision_last_steps: int = 3, + mixed_precision_reasoner_policy: str = "high_precision", ) -> Cosmos3OmniPipelineOutput: r""" Run the Cosmos 3 omni pipeline end-to-end: encode the (optional) conditioning image/video, denoise vision and @@ -1435,6 +1444,18 @@ def __call__( When `True` and a `CosmosSafetyChecker` is attached, runs the text guardrail on the prompt before generation and the video guardrail on the decoded frames. Set to `False` to skip both for this call; the checker remains loaded for subsequent calls. + mixed_precision_format (`str`, *optional*, defaults to `"none"`): + Set to `"fp8"` to use mixed W8A8/W8A16 denoising on a serialized ModelOpt FP8 checkpoint. Disabled + unless this is `"fp8"`. + mixed_precision_first_steps (`int`, *optional*, defaults to `3`): + Number of leading denoising steps that run W8A16 (activation quantizers disabled) when mixed precision + is enabled. + mixed_precision_last_steps (`int`, *optional*, defaults to `3`): + Number of trailing denoising steps that run W8A16 when mixed precision is enabled. Middle steps keep + native W8A8. Precision is selected once per scheduler step so CFG cond/uncond calls match. + mixed_precision_reasoner_policy (`str`, *optional*, defaults to `"high_precision"`): + Whether the reasoner path always uses W8A16 (`"high_precision"`) or follows the checkpoint's native + W8A8 path (`"base_precision"`). Returns: [`Cosmos3OmniPipelineOutput`] or `tuple`: @@ -1685,98 +1706,76 @@ def __call__( # 7. Denoising loop num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order self._num_timesteps = len(timesteps) - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - self._current_timestep = t - timestep = t.item() - - # The transformer projections (proj_in / audio_proj_in) are bf16; cast the per-step - # noisy tokens before packing so the modality tokens enter the model in the right dtype. - vision_tokens = latents.to(device=device, dtype=dtype) - sound_tokens = sound_latents.to(device=device, dtype=dtype) if sound_latents is not None else None - action_tokens = action_latents.to(device=device, dtype=dtype) if action_latents is not None else None - # The static packs both report the same num_noisy_vision_tokens / sound_len, so a - # single per-step timestep tensor per modality is shared by the cond / uncond passes. - vision_timesteps = torch.full((num_noisy_vision_tokens,), timestep, device=device) - sound_timesteps = ( - torch.full((sound_len,), timestep, device=device) if sound_tokens is not None else None - ) - action_timesteps = ( - torch.full((action_noisy_len,), timestep, device=device) if action_tokens is not None else None - ) - - # --- Conditional pass --- - preds_vision, preds_sound, preds_action = self.transformer( - input_ids=cond_packed_static["input_ids"], - text_indexes=cond_packed_static["text_indexes"], - position_ids=cond_packed_static["position_ids"], - und_len=cond_packed_static["und_len"], - sequence_length=cond_packed_static["sequence_length"], - vision_tokens=[vision_tokens], - vision_token_shapes=cond_packed_static["vision_token_shapes"], - vision_sequence_indexes=cond_packed_static["vision_sequence_indexes"], - vision_mse_loss_indexes=cond_packed_static["vision_mse_loss_indexes"], - vision_timesteps=vision_timesteps, - vision_noisy_frame_indexes=cond_packed_static["vision_noisy_frame_indexes"], - sound_tokens=[sound_tokens] if sound_tokens is not None else None, - sound_token_shapes=cond_packed_static.get("sound_token_shapes"), - sound_sequence_indexes=cond_packed_static.get("sound_sequence_indexes"), - sound_mse_loss_indexes=cond_packed_static.get("sound_mse_loss_indexes"), - sound_timesteps=sound_timesteps, - sound_noisy_frame_indexes=cond_packed_static.get("sound_noisy_frame_indexes"), - action_tokens=[action_tokens] if action_tokens is not None else None, - action_token_shapes=cond_packed_static.get("action_token_shapes"), - action_sequence_indexes=cond_packed_static.get("action_sequence_indexes"), - action_mse_loss_indexes=cond_packed_static.get("action_mse_loss_indexes"), - action_timesteps=action_timesteps, - action_noisy_frame_indexes=cond_packed_static.get("action_noisy_frame_indexes"), - action_domain_ids=[action_domain_id] if action_domain_id is not None else None, - return_dict=False, - ) - cond_v_vision, cond_v_sound, cond_v_action = self._mask_velocity_predictions( - preds_vision, - preds_sound, - vision_condition_mask=[vision_condition_mask], - sound_condition_mask=[sound_condition_mask] if sound_condition_mask is not None else None, - preds_action=preds_action, - action_condition_mask=[action_condition_mask] if action_condition_mask is not None else None, - raw_action_dim=raw_action_dim_resolved, - ) + mixed_precision = Cosmos3MixedPrecisionConfig.from_kwargs( + 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, + ) + self._mixed_precision_config = mixed_precision + self._mixed_precision_trace = [] + try: + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + apply_cosmos3_mixed_precision_step( + self.transformer, + mixed_precision, + i, + len(timesteps), + trace=self._mixed_precision_trace, + ) + if self.interrupt: + continue + + self._current_timestep = t + timestep = t.item() + + # The transformer projections (proj_in / audio_proj_in) are bf16; cast the per-step + # noisy tokens before packing so the modality tokens enter the model in the right dtype. + vision_tokens = latents.to(device=device, dtype=dtype) + sound_tokens = sound_latents.to(device=device, dtype=dtype) if sound_latents is not None else None + action_tokens = ( + action_latents.to(device=device, dtype=dtype) if action_latents is not None else None + ) + # The static packs both report the same num_noisy_vision_tokens / sound_len, so a + # single per-step timestep tensor per modality is shared by the cond / uncond passes. + vision_timesteps = torch.full((num_noisy_vision_tokens,), timestep, device=device) + sound_timesteps = ( + torch.full((sound_len,), timestep, device=device) if sound_tokens is not None else None + ) + action_timesteps = ( + torch.full((action_noisy_len,), timestep, device=device) if action_tokens is not None else None + ) - # --- Unconditional pass (Skip if not using CFG) --- - uncond_v_vision = uncond_v_sound = uncond_v_action = None - if self.do_classifier_free_guidance: + # --- Conditional pass --- preds_vision, preds_sound, preds_action = self.transformer( - input_ids=uncond_packed_static["input_ids"], - text_indexes=uncond_packed_static["text_indexes"], - position_ids=uncond_packed_static["position_ids"], - und_len=uncond_packed_static["und_len"], - sequence_length=uncond_packed_static["sequence_length"], + input_ids=cond_packed_static["input_ids"], + text_indexes=cond_packed_static["text_indexes"], + position_ids=cond_packed_static["position_ids"], + und_len=cond_packed_static["und_len"], + sequence_length=cond_packed_static["sequence_length"], vision_tokens=[vision_tokens], - vision_token_shapes=uncond_packed_static["vision_token_shapes"], - vision_sequence_indexes=uncond_packed_static["vision_sequence_indexes"], - vision_mse_loss_indexes=uncond_packed_static["vision_mse_loss_indexes"], + vision_token_shapes=cond_packed_static["vision_token_shapes"], + vision_sequence_indexes=cond_packed_static["vision_sequence_indexes"], + vision_mse_loss_indexes=cond_packed_static["vision_mse_loss_indexes"], vision_timesteps=vision_timesteps, - vision_noisy_frame_indexes=uncond_packed_static["vision_noisy_frame_indexes"], + vision_noisy_frame_indexes=cond_packed_static["vision_noisy_frame_indexes"], sound_tokens=[sound_tokens] if sound_tokens is not None else None, - sound_token_shapes=uncond_packed_static.get("sound_token_shapes"), - sound_sequence_indexes=uncond_packed_static.get("sound_sequence_indexes"), - sound_mse_loss_indexes=uncond_packed_static.get("sound_mse_loss_indexes"), + sound_token_shapes=cond_packed_static.get("sound_token_shapes"), + sound_sequence_indexes=cond_packed_static.get("sound_sequence_indexes"), + sound_mse_loss_indexes=cond_packed_static.get("sound_mse_loss_indexes"), sound_timesteps=sound_timesteps, - sound_noisy_frame_indexes=uncond_packed_static.get("sound_noisy_frame_indexes"), + sound_noisy_frame_indexes=cond_packed_static.get("sound_noisy_frame_indexes"), action_tokens=[action_tokens] if action_tokens is not None else None, - action_token_shapes=uncond_packed_static.get("action_token_shapes"), - action_sequence_indexes=uncond_packed_static.get("action_sequence_indexes"), - action_mse_loss_indexes=uncond_packed_static.get("action_mse_loss_indexes"), + action_token_shapes=cond_packed_static.get("action_token_shapes"), + action_sequence_indexes=cond_packed_static.get("action_sequence_indexes"), + action_mse_loss_indexes=cond_packed_static.get("action_mse_loss_indexes"), action_timesteps=action_timesteps, - action_noisy_frame_indexes=uncond_packed_static.get("action_noisy_frame_indexes"), + action_noisy_frame_indexes=cond_packed_static.get("action_noisy_frame_indexes"), action_domain_ids=[action_domain_id] if action_domain_id is not None else None, return_dict=False, ) - uncond_v_vision, uncond_v_sound, uncond_v_action = self._mask_velocity_predictions( + cond_v_vision, cond_v_sound, cond_v_action = self._mask_velocity_predictions( preds_vision, preds_sound, vision_condition_mask=[vision_condition_mask], @@ -1786,53 +1785,98 @@ def __call__( raw_action_dim=raw_action_dim_resolved, ) - # --- CFG combine + per-modality scheduler step --- - # UniPC's multistep_uni_p_bh_update einsum ("k,bkc...->bc...") requires sample - # to carry a batch dim; per-modality latents have no batch axis, so wrap for the step. - - # Skip CFG for 1.0 guidance scale - if self.do_classifier_free_guidance: - velocity_vision = uncond_v_vision + guidance_scale * (cond_v_vision - uncond_v_vision) - else: - velocity_vision = cond_v_vision - - latents = self.scheduler.step( - velocity_vision.unsqueeze(0), t, latents.unsqueeze(0), return_dict=False - )[0].squeeze(0) + # --- Unconditional pass (Skip if not using CFG) --- + uncond_v_vision = uncond_v_sound = uncond_v_action = None + if self.do_classifier_free_guidance: + preds_vision, preds_sound, preds_action = self.transformer( + input_ids=uncond_packed_static["input_ids"], + text_indexes=uncond_packed_static["text_indexes"], + position_ids=uncond_packed_static["position_ids"], + und_len=uncond_packed_static["und_len"], + sequence_length=uncond_packed_static["sequence_length"], + vision_tokens=[vision_tokens], + vision_token_shapes=uncond_packed_static["vision_token_shapes"], + vision_sequence_indexes=uncond_packed_static["vision_sequence_indexes"], + vision_mse_loss_indexes=uncond_packed_static["vision_mse_loss_indexes"], + vision_timesteps=vision_timesteps, + vision_noisy_frame_indexes=uncond_packed_static["vision_noisy_frame_indexes"], + sound_tokens=[sound_tokens] if sound_tokens is not None else None, + sound_token_shapes=uncond_packed_static.get("sound_token_shapes"), + sound_sequence_indexes=uncond_packed_static.get("sound_sequence_indexes"), + sound_mse_loss_indexes=uncond_packed_static.get("sound_mse_loss_indexes"), + sound_timesteps=sound_timesteps, + sound_noisy_frame_indexes=uncond_packed_static.get("sound_noisy_frame_indexes"), + action_tokens=[action_tokens] if action_tokens is not None else None, + action_token_shapes=uncond_packed_static.get("action_token_shapes"), + action_sequence_indexes=uncond_packed_static.get("action_sequence_indexes"), + action_mse_loss_indexes=uncond_packed_static.get("action_mse_loss_indexes"), + action_timesteps=action_timesteps, + action_noisy_frame_indexes=uncond_packed_static.get("action_noisy_frame_indexes"), + action_domain_ids=[action_domain_id] if action_domain_id is not None else None, + return_dict=False, + ) + uncond_v_vision, uncond_v_sound, uncond_v_action = self._mask_velocity_predictions( + preds_vision, + preds_sound, + vision_condition_mask=[vision_condition_mask], + sound_condition_mask=[sound_condition_mask] if sound_condition_mask is not None else None, + preds_action=preds_action, + action_condition_mask=[action_condition_mask] + if action_condition_mask is not None + else None, + raw_action_dim=raw_action_dim_resolved, + ) + + # --- CFG combine + per-modality scheduler step --- + # UniPC's multistep_uni_p_bh_update einsum ("k,bkc...->bc...") requires sample + # to carry a batch dim; per-modality latents have no batch axis, so wrap for the step. - if sound_scheduler is not None and cond_v_sound is not None: # Skip CFG for 1.0 guidance scale if self.do_classifier_free_guidance: - velocity_sound = uncond_v_sound + guidance_scale * (cond_v_sound - uncond_v_sound) + velocity_vision = uncond_v_vision + guidance_scale * (cond_v_vision - uncond_v_vision) else: - velocity_sound = cond_v_sound - sound_latents = sound_scheduler.step( - velocity_sound.unsqueeze(0), t, sound_latents.unsqueeze(0), return_dict=False - )[0].squeeze(0) + velocity_vision = cond_v_vision - has_noisy_action = ( - action_condition_mask is not None and action_condition_mask.sum() < action_condition_mask.numel() - ) - if action_scheduler is not None and has_noisy_action and cond_v_action is not None: - if self.do_classifier_free_guidance: - velocity_action = uncond_v_action + guidance_scale * (cond_v_action - uncond_v_action) - else: - velocity_action = cond_v_action - action_latents = action_scheduler.step( - velocity_action.unsqueeze(0), t, action_latents.unsqueeze(0), return_dict=False + latents = self.scheduler.step( + velocity_vision.unsqueeze(0), t, latents.unsqueeze(0), return_dict=False )[0].squeeze(0) - if raw_action_dim_resolved is not None: - action_latents[:, raw_action_dim_resolved:] = 0 - - if callback_on_step_end is not None: - callback_kwargs = {} - for key in callback_on_step_end_tensor_inputs: - callback_kwargs[key] = locals()[key] - callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) - latents = callback_outputs.pop("latents", latents) - - if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): - progress_bar.update() + + if sound_scheduler is not None and cond_v_sound is not None: + # Skip CFG for 1.0 guidance scale + if self.do_classifier_free_guidance: + velocity_sound = uncond_v_sound + guidance_scale * (cond_v_sound - uncond_v_sound) + else: + velocity_sound = cond_v_sound + sound_latents = sound_scheduler.step( + velocity_sound.unsqueeze(0), t, sound_latents.unsqueeze(0), return_dict=False + )[0].squeeze(0) + + has_noisy_action = ( + action_condition_mask is not None + and action_condition_mask.sum() < action_condition_mask.numel() + ) + if action_scheduler is not None and has_noisy_action and cond_v_action is not None: + if self.do_classifier_free_guidance: + velocity_action = uncond_v_action + guidance_scale * (cond_v_action - uncond_v_action) + else: + velocity_action = cond_v_action + action_latents = action_scheduler.step( + velocity_action.unsqueeze(0), t, action_latents.unsqueeze(0), return_dict=False + )[0].squeeze(0) + if raw_action_dim_resolved is not None: + action_latents[:, raw_action_dim_resolved:] = 0 + + if callback_on_step_end is not None: + callback_kwargs = {} + for key in callback_on_step_end_tensor_inputs: + callback_kwargs[key] = locals()[key] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + latents = callback_outputs.pop("latents", latents) + + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + finally: + reset_cosmos3_mixed_precision(self.transformer, mixed_precision) self._current_timestep = None diff --git a/tests/pipelines/cosmos/test_cosmos3_mixed_precision.py b/tests/pipelines/cosmos/test_cosmos3_mixed_precision.py new file mode 100644 index 000000000000..1134b52d2d80 --- /dev/null +++ b/tests/pipelines/cosmos/test_cosmos3_mixed_precision.py @@ -0,0 +1,160 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +from types import SimpleNamespace + +import torch +import torch.nn.functional as F + +from diffusers.pipelines.cosmos.mixed_precision import ( + Cosmos3MixedPrecisionConfig, + apply_cosmos3_mixed_precision_step, + reset_cosmos3_mixed_precision, +) + + +class _FakeModelOptFp8Linear(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter( + torch.tensor([[4.0, -2.0], [1.0, 3.0]], dtype=torch.float8_e4m3fn), + requires_grad=False, + ) + self.bias = None + self.input_quantizer = SimpleNamespace(is_enabled=True) + self.weight_quantizer = SimpleNamespace(_scale=torch.tensor(0.25), is_enabled=True) + self._should_run_real_quant_gemm = True + + def forward(self, inputs): + # Stand in for ModelOpt's native W8A8 GEMM. + quantized_inputs = inputs.to(torch.float8_e4m3fn).to(inputs.dtype) + dense_weight = self.weight.to(inputs.dtype) * self.weight_quantizer._scale.to(inputs.dtype) + return F.linear(quantized_inputs, dense_weight, self.bias) + + +class _FakeAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.to_q = _FakeModelOptFp8Linear() + self.add_q_proj = _FakeModelOptFp8Linear() + + +class _FakeLayer(torch.nn.Module): + def __init__(self): + super().__init__() + self.self_attn = _FakeAttention() + + +class _Transformer(torch.nn.Module): + def __init__(self): + super().__init__() + self.layers = torch.nn.ModuleList([_FakeLayer()]) + + +class Cosmos3MixedPrecisionConfigTests(unittest.TestCase): + def test_default_50_step_schedule_matches_vllm_omni(self): + config = Cosmos3MixedPrecisionConfig(format="fp8") + selected = [i for i in range(50) if config.use_high_precision(i, 50)] + self.assertEqual(selected, [0, 1, 2, 47, 48, 49]) + self.assertEqual(config.precision_name(0, 50), "W8A16") + self.assertEqual(config.precision_name(25, 50), "W8A8") + + def test_asymmetric_and_overlap_boundaries(self): + cases = [ + (2, 4, 10, [0, 1, 6, 7, 8, 9]), + (0, 2, 7, [5, 6]), + (2, 0, 7, [0, 1]), + (0, 0, 7, []), + (4, 4, 7, list(range(7))), + ] + for first_steps, last_steps, num_steps, selected in cases: + with self.subTest(first=first_steps, last=last_steps, n=num_steps): + config = Cosmos3MixedPrecisionConfig(format="fp8", first_steps=first_steps, last_steps=last_steps) + self.assertEqual( + [i for i in range(num_steps) if config.use_high_precision(i, num_steps)], + selected, + ) + + def test_one_step_keeps_base_precision(self): + config = Cosmos3MixedPrecisionConfig(format="fp8", first_steps=1, last_steps=1) + self.assertFalse(config.use_high_precision(0, 1)) + + def test_disabled_format_is_noop(self): + config = Cosmos3MixedPrecisionConfig.from_kwargs(mixed_precision_format="none") + self.assertFalse(config.enabled) + transformer = _Transformer() + trace = [] + inputs = torch.tensor([[1.1, -0.7]], dtype=torch.bfloat16) + expected = transformer.layers[0].self_attn.to_q(inputs) + name = apply_cosmos3_mixed_precision_step(transformer, config, 0, 10, trace=trace) + self.assertEqual(name, "base") + torch.testing.assert_close(transformer.layers[0].self_attn.to_q(inputs), expected) + + def test_rejects_invalid_values(self): + with self.assertRaises(ValueError): + Cosmos3MixedPrecisionConfig.from_kwargs(mixed_precision_format="nvfp4") + with self.assertRaises(ValueError): + Cosmos3MixedPrecisionConfig.from_kwargs(mixed_precision_reasoner_policy="fp16") + with self.assertRaises(TypeError): + Cosmos3MixedPrecisionConfig.from_kwargs(mixed_precision_first_steps=-1) + + def test_dispatches_generation_w8a16_edges_and_w8a8_middle(self): + config = Cosmos3MixedPrecisionConfig(format="fp8", first_steps=1, last_steps=1) + transformer = _Transformer() + generation = transformer.layers[0].self_attn.to_q + reasoner = transformer.layers[0].self_attn.add_q_proj + inputs = torch.tensor([[1.1, -0.7]], dtype=torch.bfloat16) + expected_w8a8 = generation(inputs) + expected_w8a16 = F.linear( + inputs, + generation.weight.to(inputs.dtype) * generation.weight_quantizer._scale.to(inputs.dtype), + ) + self.assertFalse(torch.equal(expected_w8a8, expected_w8a16)) + + trace = [] + apply_cosmos3_mixed_precision_step(transformer, config, 0, 5, trace=trace) + torch.testing.assert_close(generation(inputs), expected_w8a16) + torch.testing.assert_close(reasoner(inputs), expected_w8a16) + + apply_cosmos3_mixed_precision_step(transformer, config, 2, 5, trace=trace) + torch.testing.assert_close(generation(inputs), expected_w8a8) + # The default policy keeps the reasoner in high precision on every step. + torch.testing.assert_close(reasoner(inputs), expected_w8a16) + + apply_cosmos3_mixed_precision_step(transformer, config, 4, 5, trace=trace) + torch.testing.assert_close(generation(inputs), expected_w8a16) + self.assertEqual(trace, ["W8A16", "W8A8", "W8A16"]) + + reset_cosmos3_mixed_precision(transformer, config) + torch.testing.assert_close(generation(inputs), expected_w8a8) + torch.testing.assert_close(reasoner(inputs), expected_w8a8) + + def test_reasoner_can_use_base_precision(self): + config = Cosmos3MixedPrecisionConfig( + format="fp8", + first_steps=0, + last_steps=0, + reasoner_policy="base_precision", + ) + transformer = _Transformer() + reasoner = transformer.layers[0].self_attn.add_q_proj + inputs = torch.tensor([[1.1, -0.7]], dtype=torch.bfloat16) + expected_w8a8 = reasoner(inputs) + apply_cosmos3_mixed_precision_step(transformer, config, 2, 5) + torch.testing.assert_close(reasoner(inputs), expected_w8a8) + + +if __name__ == "__main__": + unittest.main() From 37c9248c326f62103baaf0fa26eee5649afb27b3 Mon Sep 17 00:00:00 2001 From: Yiming Zhao Date: Mon, 31 Aug 2026 11:36:15 -0700 Subject: [PATCH 2/3] Read Cosmos3 mixed-precision schedule from the checkpoint policy. 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. --- .../modular_pipelines/cosmos/denoise.py | 30 +- .../pipelines/cosmos/mixed_precision.py | 278 +++++++++++++++++- .../pipelines/cosmos/pipeline_cosmos3_omni.py | 37 +-- .../cosmos/test_cosmos3_mixed_precision.py | 71 +++++ 4 files changed, 371 insertions(+), 45 deletions(-) diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py index c7003e593dbb..7c5327f0ab53 100644 --- a/src/diffusers/modular_pipelines/cosmos/denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/denoise.py @@ -471,40 +471,38 @@ def loop_inputs(self) -> list[InputParam]: InputParam( name="mixed_precision_format", type_hint=str, - default="none", - description="Set to 'fp8' to enable mixed W8A8/W8A16 denoising.", + default=None, + description="None reads the checkpoint diffusion_step_policy; 'fp8' forces mixed precision; 'none' disables it.", ), InputParam( name="mixed_precision_first_steps", type_hint=int, - default=3, - description="Leading W8A16 step count when mixed precision is enabled.", + default=None, + description="Optional override for the leading W8A16 step count.", ), InputParam( name="mixed_precision_last_steps", type_hint=int, - default=3, - description="Trailing W8A16 step count when mixed precision is enabled.", + default=None, + description="Optional override for the trailing W8A16 step count.", ), InputParam( name="mixed_precision_reasoner_policy", type_hint=str, - default="high_precision", - description="Use W8A16 or checkpoint-native W8A8 for the reasoner path.", + default=None, + description="Optional override: W8A16 ('high_precision') or native W8A8 ('base_precision') for the reasoner path.", ), ] @torch.no_grad() def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - first_steps = getattr(block_state, "mixed_precision_first_steps", None) - last_steps = getattr(block_state, "mixed_precision_last_steps", None) - mixed_precision = Cosmos3MixedPrecisionConfig.from_kwargs( - mixed_precision_format=getattr(block_state, "mixed_precision_format", None) or "none", - mixed_precision_first_steps=3 if first_steps is None else first_steps, - mixed_precision_last_steps=3 if last_steps is None else last_steps, - mixed_precision_reasoner_policy=getattr(block_state, "mixed_precision_reasoner_policy", None) - or "high_precision", + mixed_precision = Cosmos3MixedPrecisionConfig.resolve( + components.transformer, + mixed_precision_format=getattr(block_state, "mixed_precision_format", None), + mixed_precision_first_steps=getattr(block_state, "mixed_precision_first_steps", None), + mixed_precision_last_steps=getattr(block_state, "mixed_precision_last_steps", None), + mixed_precision_reasoner_policy=getattr(block_state, "mixed_precision_reasoner_policy", None), ) trace = [] try: diff --git a/src/diffusers/pipelines/cosmos/mixed_precision.py b/src/diffusers/pipelines/cosmos/mixed_precision.py index 7e641d9010da..6e6f41b26d4a 100644 --- a/src/diffusers/pipelines/cosmos/mixed_precision.py +++ b/src/diffusers/pipelines/cosmos/mixed_precision.py @@ -18,6 +18,10 @@ bypasses that forward, dequantizes the FP8 weight to the activation dtype, and uses :func:`torch.nn.functional.linear`. This matches the uncached strategy in vLLM-Omni (vllm-project/vllm-omni#6560). + +Schedule defaults come from ``quantization_config.runtime.diffusion_step_policy`` +on the transformer (Cosmos3-Experimental discussion #19). Distilled checkpoints +omit that policy and stay on native W8A8. """ from __future__ import annotations @@ -33,9 +37,27 @@ MixedPrecisionFormat = Literal["none", "fp8"] ReasonerPolicy = Literal["high_precision", "base_precision"] +OverlapMode = Literal["a16", "native"] MIXED_PRECISION_FORMATS = frozenset({"none", "fp8"}) REASONER_POLICIES = frozenset({"high_precision", "base_precision"}) +OVERLAP_MODES = frozenset({"a16", "native"}) _RUNTIME_ATTRIBUTE = "_cosmos3_mixed_precision_runtime" +_SUPPORTED_POLICY_TYPE = "first_last_n" +_SUPPORTED_INDEX_SPACE = "denoising_loop_iteration" +_SUPPORTED_SCHEMA_VERSION = 1 + +# Official Nano/Super/Super-I2V FP8 policy (nvidia/Cosmos3-Experimental#19). +FIRST_LAST_N_FP8_POLICY = { + "schema_version": 1, + "type": "first_last_n", + "index_space": "denoising_loop_iteration", + "scope": ["transformer"], + "default_mode": "native", + "first_steps": {"count": 3, "mode": "a16"}, + "last_steps": {"count": 3, "mode": "a16"}, + "overlap": "a16", + "reasoner": "a16", +} @dataclass(frozen=True) @@ -46,6 +68,7 @@ class Cosmos3MixedPrecisionConfig: first_steps: int = 3 last_steps: int = 3 reasoner_policy: ReasonerPolicy = "high_precision" + overlap: OverlapMode = "a16" @classmethod def from_kwargs( @@ -54,24 +77,81 @@ def from_kwargs( mixed_precision_first_steps: int = 3, mixed_precision_last_steps: int = 3, mixed_precision_reasoner_policy: str = "high_precision", + mixed_precision_overlap: str = "a16", + ) -> Cosmos3MixedPrecisionConfig: + return cls.resolve( + 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, + mixed_precision_overlap=mixed_precision_overlap, + ) + + @classmethod + def resolve( + cls, + transformer: nn.Module | None = None, + *, + mixed_precision_format: str | None = None, + mixed_precision_first_steps: int | None = None, + mixed_precision_last_steps: int | None = None, + mixed_precision_reasoner_policy: str | None = None, + mixed_precision_overlap: str | None = None, + quantization_config: dict[str, Any] | None = None, ) -> Cosmos3MixedPrecisionConfig: - precision_format = str(mixed_precision_format).strip().lower() - if precision_format not in MIXED_PRECISION_FORMATS: + """Build a schedule from the checkpoint policy, with optional call-site overrides. + + ``mixed_precision_format=None`` (pipeline default) means auto: enable mixed + precision only when the transformer declares ``diffusion_step_policy``. + Distilled FP8 checkpoints omit that field and stay native W8A8. Pass + ``"fp8"`` to force the schedule, or ``"none"`` to disable it. + """ + parsed = _parsed_checkpoint_policy(transformer, quantization_config) + format_override = _optional_lower(mixed_precision_format) + if format_override is not None and format_override not in MIXED_PRECISION_FORMATS: raise ValueError( "mixed_precision_format must be one of " - f"{sorted(MIXED_PRECISION_FORMATS)}, got {mixed_precision_format!r}" + f"{sorted(MIXED_PRECISION_FORMATS)} or None, got {mixed_precision_format!r}" ) - reasoner_policy = str(mixed_precision_reasoner_policy).strip().lower() - if reasoner_policy not in REASONER_POLICIES: - raise ValueError( - "mixed_precision_reasoner_policy must be one of " - f"{sorted(REASONER_POLICIES)}, got {mixed_precision_reasoner_policy!r}" + if mixed_precision_first_steps is not None: + mixed_precision_first_steps = _non_negative_int( + mixed_precision_first_steps, "mixed_precision_first_steps" + ) + if mixed_precision_last_steps is not None: + mixed_precision_last_steps = _non_negative_int( + mixed_precision_last_steps, "mixed_precision_last_steps" ) + if mixed_precision_reasoner_policy is not None: + mixed_precision_reasoner_policy = _validated_reasoner(mixed_precision_reasoner_policy) + if mixed_precision_overlap is not None: + mixed_precision_overlap = _validated_overlap(mixed_precision_overlap) + + if format_override == "none": + return cls(format="none") + + if parsed is None and format_override != "fp8": + return cls(format="none") + + first_steps = parsed.first_steps if parsed is not None else 3 + last_steps = parsed.last_steps if parsed is not None else 3 + reasoner_policy = parsed.reasoner_policy if parsed is not None else "high_precision" + overlap = parsed.overlap if parsed is not None else "a16" + + if mixed_precision_first_steps is not None: + first_steps = mixed_precision_first_steps + if mixed_precision_last_steps is not None: + last_steps = mixed_precision_last_steps + if mixed_precision_reasoner_policy is not None: + reasoner_policy = mixed_precision_reasoner_policy + if mixed_precision_overlap is not None: + overlap = mixed_precision_overlap + return cls( - format=precision_format, # type: ignore[arg-type] - first_steps=_non_negative_int(mixed_precision_first_steps, "mixed_precision_first_steps"), - last_steps=_non_negative_int(mixed_precision_last_steps, "mixed_precision_last_steps"), + format="fp8", + first_steps=first_steps, + last_steps=last_steps, reasoner_policy=reasoner_policy, # type: ignore[arg-type] + overlap=overlap, # type: ignore[arg-type] ) @property @@ -87,7 +167,11 @@ def use_high_precision(self, step_index: int, num_steps: int) -> bool: # Match vLLM-Omni: a 1-step request keeps the checkpoint's base (W8A8) path. if num_steps == 1: return False - return step_index < self.first_steps or step_index >= num_steps - self.last_steps + in_first = step_index < self.first_steps + in_last = step_index >= num_steps - self.last_steps + if in_first and in_last: + return self.overlap == "a16" + return in_first or in_last def precision_name(self, step_index: int, num_steps: int) -> str: if not self.enabled: @@ -253,3 +337,173 @@ def _non_negative_int(value: Any, name: str) -> int: if not isinstance(value, int) or isinstance(value, bool) or value < 0: raise TypeError(f"{name} must be a non-negative integer, got {value!r}") return value + + +@dataclass(frozen=True) +class _ParsedCheckpointPolicy: + first_steps: int + last_steps: int + reasoner_policy: ReasonerPolicy + overlap: OverlapMode + + +def _optional_lower(value: str | None) -> str | None: + if value is None: + return None + return str(value).strip().lower() + + +def _validated_reasoner(value: str) -> ReasonerPolicy: + reasoner_policy = str(value).strip().lower() + if reasoner_policy not in REASONER_POLICIES: + raise ValueError( + "mixed_precision_reasoner_policy must be one of " + f"{sorted(REASONER_POLICIES)}, got {value!r}" + ) + return reasoner_policy # type: ignore[return-value] + + +def _validated_overlap(value: str) -> OverlapMode: + overlap = str(value).strip().lower() + if overlap not in OVERLAP_MODES: + raise ValueError(f"overlap must be one of {sorted(OVERLAP_MODES)}, got {value!r}") + return overlap # type: ignore[return-value] + + +def _maybe_mapping(value: Any) -> dict[str, Any] | None: + if value is None: + return None + if hasattr(value, "to_dict"): + value = value.to_dict() + if isinstance(value, dict): + return dict(value) + return None + + +def quantization_config_from_module(module: nn.Module | None) -> dict[str, Any] | None: + """Read ModelOpt ``quantization_config`` from a loaded transformer, if present.""" + if module is None: + return None + candidates: list[Any] = [getattr(module, "quantization_config", None)] + config = getattr(module, "config", None) + if config is not None: + candidates.append(getattr(config, "quantization_config", None)) + if hasattr(config, "get"): + candidates.append(config.get("quantization_config")) + for candidate in candidates: + mapped = _maybe_mapping(candidate) + if mapped is not None: + return mapped + + 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 name_or_path and callable(loader): + 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): + mapped = _maybe_mapping(raw.get("quantization_config")) + if mapped is not None: + return mapped + return None + + +def _parsed_checkpoint_policy( + transformer: nn.Module | None, + quantization_config: dict[str, Any] | None, +) -> _ParsedCheckpointPolicy | None: + quant_config = _maybe_mapping(quantization_config) or quantization_config_from_module(transformer) + if not quant_config: + return None + runtime = quant_config.get("runtime") + if runtime is None: + return None + runtime_map = _maybe_mapping(runtime) + if runtime_map is None: + raise ValueError("quantization_config.runtime must be a mapping or null") + policy = runtime_map.get("diffusion_step_policy") + if policy is None: + return None + return parse_diffusion_step_policy(policy, quant_config) + + +def parse_diffusion_step_policy( + policy: Any, + quantization_config: dict[str, Any] | None = None, +) -> _ParsedCheckpointPolicy: + """Validate the versioned first/last-N policy from transformer/config.json.""" + policy_map = _maybe_mapping(policy) + if policy_map is None: + raise ValueError("diffusion_step_policy must be a mapping") + + schema_version = policy_map.get("schema_version", _SUPPORTED_SCHEMA_VERSION) + if schema_version != _SUPPORTED_SCHEMA_VERSION: + raise ValueError( + f"Unsupported diffusion_step_policy.schema_version={schema_version}; " + f"Diffusers supports version {_SUPPORTED_SCHEMA_VERSION}" + ) + policy_type = policy_map.get("type") + if policy_type != _SUPPORTED_POLICY_TYPE: + raise ValueError( + f"Unsupported diffusion_step_policy.type={policy_type!r}; " + f"expected {_SUPPORTED_POLICY_TYPE!r}" + ) + index_space = policy_map.get("index_space", _SUPPORTED_INDEX_SPACE) + if index_space != _SUPPORTED_INDEX_SPACE: + raise ValueError( + f"Unsupported diffusion_step_policy.index_space={index_space!r}; " + f"expected {_SUPPORTED_INDEX_SPACE!r}" + ) + + if quantization_config: + algo = str( + quantization_config.get("quant_algo") + or quantization_config.get("quant_type") + or "" + ).upper() + if algo and "FP8" not in algo: + raise ValueError( + "Cosmos3 mixed precision in Diffusers currently supports ModelOpt FP8 " + f"checkpoints, got quant_algo/quant_type={algo!r}" + ) + + first_steps = _window_count(policy_map.get("first_steps"), "first_steps") + last_steps = _window_count(policy_map.get("last_steps"), "last_steps") + overlap = _validated_overlap(str(policy_map.get("overlap", "a16"))) + reasoner = policy_map.get("reasoner", "a16") + if reasoner == "a16": + reasoner_policy: ReasonerPolicy = "high_precision" + elif reasoner == "native": + reasoner_policy = "base_precision" + else: + raise ValueError(f"diffusion_step_policy.reasoner must be 'a16' or 'native', got {reasoner!r}") + return _ParsedCheckpointPolicy( + first_steps=first_steps, + last_steps=last_steps, + reasoner_policy=reasoner_policy, + overlap=overlap, + ) + + +def _window_count(spec: Any, name: str) -> int: + spec_map = _maybe_mapping(spec) + if spec_map is None: + raise ValueError(f"diffusion_step_policy.{name} must be a mapping with count and mode") + mode = spec_map.get("mode") + if mode not in {"a16", "native"}: + raise ValueError(f"diffusion_step_policy.{name}.mode must be 'a16' or 'native', got {mode!r}") + count = _non_negative_int(spec_map.get("count"), f"diffusion_step_policy.{name}.count") + if mode == "native": + return 0 + return count + diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py index 09309105f60c..0fa6b1cf2be1 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py @@ -1345,10 +1345,10 @@ def __call__( add_resolution_template: bool = True, add_duration_template: bool = True, enable_safety_check: bool = True, - mixed_precision_format: str = "none", - mixed_precision_first_steps: int = 3, - mixed_precision_last_steps: int = 3, - mixed_precision_reasoner_policy: str = "high_precision", + mixed_precision_format: str | None = None, + mixed_precision_first_steps: int | None = None, + mixed_precision_last_steps: int | None = None, + mixed_precision_reasoner_policy: str | None = None, ) -> Cosmos3OmniPipelineOutput: r""" Run the Cosmos 3 omni pipeline end-to-end: encode the (optional) conditioning image/video, denoise vision and @@ -1444,18 +1444,20 @@ def __call__( When `True` and a `CosmosSafetyChecker` is attached, runs the text guardrail on the prompt before generation and the video guardrail on the decoded frames. Set to `False` to skip both for this call; the checker remains loaded for subsequent calls. - mixed_precision_format (`str`, *optional*, defaults to `"none"`): - Set to `"fp8"` to use mixed W8A8/W8A16 denoising on a serialized ModelOpt FP8 checkpoint. Disabled - unless this is `"fp8"`. - mixed_precision_first_steps (`int`, *optional*, defaults to `3`): - Number of leading denoising steps that run W8A16 (activation quantizers disabled) when mixed precision - is enabled. - mixed_precision_last_steps (`int`, *optional*, defaults to `3`): - Number of trailing denoising steps that run W8A16 when mixed precision is enabled. Middle steps keep - native W8A8. Precision is selected once per scheduler step so CFG cond/uncond calls match. - mixed_precision_reasoner_policy (`str`, *optional*, defaults to `"high_precision"`): - Whether the reasoner path always uses W8A16 (`"high_precision"`) or follows the checkpoint's native - W8A8 path (`"base_precision"`). + mixed_precision_format (`str`, *optional*): + `"fp8"` forces mixed W8A8/W8A16 denoising. `"none"` disables it even if the checkpoint declares a + policy. The default `None` reads `quantization_config.runtime.diffusion_step_policy` from the + transformer: Nano/Super/Super-I2V FP8 enable first/last-3 W8A16 automatically; distilled FP8 + checkpoints omit that policy and stay native W8A8. + mixed_precision_first_steps (`int`, *optional*): + Override the leading W8A16 step count from the checkpoint policy. Unused when mixed precision is + disabled. + mixed_precision_last_steps (`int`, *optional*): + Override the trailing W8A16 step count. Middle steps keep native W8A8. Precision is selected once + per scheduler step so CFG cond/uncond calls match. + mixed_precision_reasoner_policy (`str`, *optional*): + Override reasoner precision: W8A16 (`"high_precision"`) or checkpoint-native W8A8 + (`"base_precision"`). Defaults to the checkpoint policy (`reasoner: a16`). Returns: [`Cosmos3OmniPipelineOutput`] or `tuple`: @@ -1706,7 +1708,8 @@ def __call__( # 7. Denoising loop num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order self._num_timesteps = len(timesteps) - mixed_precision = Cosmos3MixedPrecisionConfig.from_kwargs( + 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, diff --git a/tests/pipelines/cosmos/test_cosmos3_mixed_precision.py b/tests/pipelines/cosmos/test_cosmos3_mixed_precision.py index 1134b52d2d80..def70d4bc44b 100644 --- a/tests/pipelines/cosmos/test_cosmos3_mixed_precision.py +++ b/tests/pipelines/cosmos/test_cosmos3_mixed_precision.py @@ -19,8 +19,10 @@ import torch.nn.functional as F from diffusers.pipelines.cosmos.mixed_precision import ( + FIRST_LAST_N_FP8_POLICY, Cosmos3MixedPrecisionConfig, apply_cosmos3_mixed_precision_step, + parse_diffusion_step_policy, reset_cosmos3_mixed_precision, ) @@ -156,5 +158,74 @@ def test_reasoner_can_use_base_precision(self): torch.testing.assert_close(reasoner(inputs), expected_w8a8) +class Cosmos3CheckpointPolicyTests(unittest.TestCase): + def _transformer_with_quant_config(self, quantization_config): + transformer = _Transformer() + transformer.config = {"quantization_config": quantization_config} + return transformer + + def test_official_policy_enables_first_last_three(self): + parsed = parse_diffusion_step_policy(FIRST_LAST_N_FP8_POLICY) + self.assertEqual(parsed.first_steps, 3) + self.assertEqual(parsed.last_steps, 3) + self.assertEqual(parsed.reasoner_policy, "high_precision") + self.assertEqual(parsed.overlap, "a16") + + transformer = self._transformer_with_quant_config( + {"quant_method": "modelopt", "quant_algo": "FP8", "runtime": {"diffusion_step_policy": FIRST_LAST_N_FP8_POLICY}} + ) + config = Cosmos3MixedPrecisionConfig.resolve(transformer) + self.assertTrue(config.enabled) + self.assertEqual([i for i in range(50) if config.use_high_precision(i, 50)], [0, 1, 2, 47, 48, 49]) + + def test_missing_runtime_stays_native_including_four_step_overlap(self): + transformer = self._transformer_with_quant_config( + {"quant_method": "modelopt", "quant_algo": "FP8", "runtime": None} + ) + config = Cosmos3MixedPrecisionConfig.resolve(transformer) + self.assertFalse(config.enabled) + # Distill 4-step must not inherit a hardcoded first/last-3 schedule. + self.assertEqual([config.precision_name(i, 4) for i in range(4)], ["base"] * 4) + + def test_explicit_none_disables_checkpoint_policy(self): + transformer = self._transformer_with_quant_config( + {"quant_method": "modelopt", "quant_algo": "FP8", "runtime": {"diffusion_step_policy": FIRST_LAST_N_FP8_POLICY}} + ) + config = Cosmos3MixedPrecisionConfig.resolve(transformer, mixed_precision_format="none") + self.assertFalse(config.enabled) + + def test_call_site_overrides_checkpoint_counts(self): + transformer = self._transformer_with_quant_config( + {"quant_method": "modelopt", "quant_algo": "FP8", "runtime": {"diffusion_step_policy": FIRST_LAST_N_FP8_POLICY}} + ) + config = Cosmos3MixedPrecisionConfig.resolve( + transformer, + mixed_precision_first_steps=1, + mixed_precision_last_steps=1, + ) + self.assertEqual([i for i in range(5) if config.use_high_precision(i, 5)], [0, 4]) + + def test_force_fp8_without_policy_uses_explicit_schedule(self): + config = Cosmos3MixedPrecisionConfig.resolve( + _Transformer(), + mixed_precision_format="fp8", + mixed_precision_first_steps=1, + mixed_precision_last_steps=1, + ) + self.assertTrue(config.enabled) + self.assertEqual([i for i in range(5) if config.use_high_precision(i, 5)], [0, 4]) + + def test_malformed_policy_fails_closed(self): + with self.assertRaises(ValueError): + parse_diffusion_step_policy({"schema_version": 2, "type": "first_last_n"}) + with self.assertRaises(ValueError): + Cosmos3MixedPrecisionConfig.resolve( + quantization_config={ + "quant_algo": "NVFP4", + "runtime": {"diffusion_step_policy": FIRST_LAST_N_FP8_POLICY}, + } + ) + + if __name__ == "__main__": unittest.main() From 6939c0508b7d1b9e2f18addbbd07f8b89b7ae098 Mon Sep 17 00:00:00 2001 From: Yiming Zhao Date: Mon, 31 Aug 2026 13:43:38 -0700 Subject: [PATCH 3/3] Harden Cosmos3 mixed-precision loading and document official Hub fp8 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. --- docs/source/en/api/pipelines/cosmos3.md | 30 ++++ .../pipelines/cosmos/mixed_precision.py | 140 ++++++++++++------ .../cosmos/test_cosmos3_mixed_precision.py | 87 ++++++++++- 3 files changed, 206 insertions(+), 51 deletions(-) diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md index ae8de9ca66c8..078a4013e443 100644 --- a/docs/source/en/api/pipelines/cosmos3.md +++ b/docs/source/en/api/pipelines/cosmos3.md @@ -43,6 +43,33 @@ Two checkpoints are released on the Hub — [`nvidia/Cosmos3-Nano`](https://hugg > [!TIP] > Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reusing-models-in-multiple-pipelines) section to learn how to efficiently load the same components into multiple pipelines. +## FP8 mixed W8A8/W8A16 denoising + +Official ModelOpt FP8 checkpoints live on the Hub `fp8` revision (for example [`nvidia/Cosmos3-Nano`](https://huggingface.co/nvidia/Cosmos3-Nano) with `revision="fp8"`). The serialized weights are static W8A8. Video Nano / Super / Super-I2V checkpoints also store a `quantization_config.runtime.diffusion_step_policy` on the transformer: the **first 3 and last 3** denoising steps run **W8A16** (dequantized FP8 weights, `torch.nn.functional.linear`), and the middle steps keep native **W8A8**. Precision is chosen once per scheduler step so CFG cond/uncond calls match. Distilled 4-step and Super-T2I FP8 checkpoints omit that policy (`runtime` is `null`) and stay native W8A8 on every step. + +Loading those weights still uses [`NVIDIAModelOptConfig`](../../quantization/modelopt) as in the ModelOpt guide. After restore, mixed precision is **on by default** when the checkpoint declares the policy — you do not pass a format flag: + +```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) # 3×W8A16 / 29×W8A8 / 3×W8A16 +``` + +Call-site overrides: + +- `mixed_precision_format="none"` disables the schedule only (quantized W8A8 remains). +- `mixed_precision_format="fp8"` forces the first/last-N schedule even if the checkpoint has no policy. +- `mixed_precision_first_steps` / `mixed_precision_last_steps` / `mixed_precision_reasoner_policy` override the checkpoint counts and reasoner path (`"high_precision"` = W8A16, `"base_precision"` = native W8A8). + +These kwargs are not Accelerate `mixed_precision`. They only select W8A8 vs W8A16 on Cosmos3 ModelOpt linears. + ## Prompt upsampling Cosmos 3 was trained on long, highly descriptive captions. For optimal quality, short text prompts should be **upsampled into a specific JSON structure** before they are passed to the pipeline. The upsampler lives in the [cosmos-framework](https://github.com/NVIDIA/cosmos-framework) package. @@ -1117,6 +1144,9 @@ config (from the checkpoint's `modular_model_index.json`) and `guidance_scale` i 1.0 since guidance is baked into the weights — passing any other value for either raises an error, and `negative_prompt` is warned about and ignored. +FP8 distilled checkpoints (`revision="fp8"`) do not declare a mixed-precision policy, so every +step stays native W8A8. + Prompts follow the same descriptive JSON structure as the non-distilled models, so short text must be upsampled first — use `--mode text2image` (T2I) or `--mode image2video` (I2V) as described in [Prompt upsampling](#prompt-upsampling), then pass the JSON via `json.dumps(...)`. diff --git a/src/diffusers/pipelines/cosmos/mixed_precision.py b/src/diffusers/pipelines/cosmos/mixed_precision.py index 6e6f41b26d4a..c32489d9960c 100644 --- a/src/diffusers/pipelines/cosmos/mixed_precision.py +++ b/src/diffusers/pipelines/cosmos/mixed_precision.py @@ -20,8 +20,8 @@ vLLM-Omni (vllm-project/vllm-omni#6560). Schedule defaults come from ``quantization_config.runtime.diffusion_step_policy`` -on the transformer (Cosmos3-Experimental discussion #19). Distilled checkpoints -omit that policy and stay on native W8A8. +in the transformer's ``config.json`` (official Hub ``revision=fp8``). Distilled +checkpoints omit that policy and stay on native W8A8. """ from __future__ import annotations @@ -45,8 +45,20 @@ _SUPPORTED_POLICY_TYPE = "first_last_n" _SUPPORTED_INDEX_SPACE = "denoising_loop_iteration" _SUPPORTED_SCHEMA_VERSION = 1 - -# Official Nano/Super/Super-I2V FP8 policy (nvidia/Cosmos3-Experimental#19). +_SUPPORTED_DEFAULT_MODE = "native" +_REQUIRED_POLICY_FIELDS = ( + "schema_version", + "type", + "index_space", + "scope", + "default_mode", + "first_steps", + "last_steps", + "overlap", + "reasoner", +) + +# Official Nano / Super / Super-I2V FP8 policy on Hub ``revision=fp8``. FIRST_LAST_N_FP8_POLICY = { "schema_version": 1, "type": "first_last_n", @@ -114,13 +126,9 @@ def resolve( f"{sorted(MIXED_PRECISION_FORMATS)} or None, got {mixed_precision_format!r}" ) if mixed_precision_first_steps is not None: - mixed_precision_first_steps = _non_negative_int( - mixed_precision_first_steps, "mixed_precision_first_steps" - ) + mixed_precision_first_steps = _non_negative_int(mixed_precision_first_steps, "mixed_precision_first_steps") if mixed_precision_last_steps is not None: - mixed_precision_last_steps = _non_negative_int( - mixed_precision_last_steps, "mixed_precision_last_steps" - ) + mixed_precision_last_steps = _non_negative_int(mixed_precision_last_steps, "mixed_precision_last_steps") if mixed_precision_reasoner_policy is not None: mixed_precision_reasoner_policy = _validated_reasoner(mixed_precision_reasoner_policy) if mixed_precision_overlap is not None: @@ -287,8 +295,8 @@ def _validate_modelopt_fp8_linear(name: str, layer: nn.Module) -> None: def _w8a16_linear(layer: nn.Module, inputs: torch.Tensor, name: str) -> torch.Tensor: - if inputs.dtype not in (torch.bfloat16, torch.float16): - raise TypeError(f"{name} W8A16 requires BF16/FP16 activations, got {inputs.dtype}") + if inputs.dtype not in (torch.bfloat16, torch.float16, torch.float32): + raise TypeError(f"{name} W8A16 requires BF16/FP16/FP32 activations, got {inputs.dtype}") scale = layer.weight_quantizer._scale.to(device=layer.weight.device, dtype=inputs.dtype) dense_weight = layer.weight.to(dtype=inputs.dtype) * scale return F.linear(inputs, dense_weight, layer.bias) @@ -325,7 +333,12 @@ def apply_cosmos3_mixed_precision_step( def reset_cosmos3_mixed_precision(module: nn.Module, config: Cosmos3MixedPrecisionConfig) -> None: - """Return installed wrappers to the checkpoint's native W8A8 path.""" + """Return installed wrappers to the checkpoint's native W8A8 path. + + Wrappers stay on ``layer.forward`` for the rest of the process. When inactive + they call the original ModelOpt forward, so a later bare ``transformer(...)`` + still runs native W8A8. + """ if not config.enabled: return runtime = getattr(module, _RUNTIME_ATTRIBUTE, None) @@ -356,10 +369,7 @@ def _optional_lower(value: str | None) -> str | None: def _validated_reasoner(value: str) -> ReasonerPolicy: reasoner_policy = str(value).strip().lower() if reasoner_policy not in REASONER_POLICIES: - raise ValueError( - "mixed_precision_reasoner_policy must be one of " - f"{sorted(REASONER_POLICIES)}, got {value!r}" - ) + raise ValueError(f"mixed_precision_reasoner_policy must be one of {sorted(REASONER_POLICIES)}, got {value!r}") return reasoner_policy # type: ignore[return-value] @@ -381,9 +391,26 @@ def _maybe_mapping(value: Any) -> dict[str, Any] | None: def quantization_config_from_module(module: nn.Module | None) -> dict[str, Any] | None: - """Read ModelOpt ``quantization_config`` from a loaded transformer, if present.""" + """Read ModelOpt ``quantization_config`` from a loaded transformer, if present. + + Diffusers strips ``quantization_config`` from the in-memory FrozenDict, and ModelOpt + restore may attach a live config without ``runtime``. Prefer the live object, then + overlay ``runtime`` from the on-disk ``transformer/config.json`` when present. + """ if module is None: return None + live = _live_quantization_config(module) + on_disk = _on_disk_quantization_config(module) + if live is None: + return on_disk + if live.get("runtime") is None and on_disk is not None and on_disk.get("runtime") is not None: + merged = dict(live) + merged["runtime"] = on_disk["runtime"] + return merged + return live + + +def _live_quantization_config(module: nn.Module) -> dict[str, Any] | None: candidates: list[Any] = [getattr(module, "quantization_config", None)] config = getattr(module, "config", None) if config is not None: @@ -394,27 +421,30 @@ def quantization_config_from_module(module: nn.Module | None) -> dict[str, Any] mapped = _maybe_mapping(candidate) if mapped is not None: return mapped + return None + +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 name_or_path and callable(loader): + 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, local_files_only=True) - except TypeError: - try: - raw = loader(name_or_path) - except Exception: - raw = None + raw = loader(name_or_path) except Exception: raw = None - if isinstance(raw, dict): - mapped = _maybe_mapping(raw.get("quantization_config")) - if mapped is not None: - return mapped + except Exception: + raw = None + if isinstance(raw, dict): + return _maybe_mapping(raw.get("quantization_config")) return None @@ -441,46 +471,63 @@ def parse_diffusion_step_policy( policy: Any, quantization_config: dict[str, Any] | None = None, ) -> _ParsedCheckpointPolicy: - """Validate the versioned first/last-N policy from transformer/config.json.""" + """Validate the versioned first/last-N policy from transformer/config.json. + + Missing fields fail closed. Official Hub policies include every key in + ``_REQUIRED_POLICY_FIELDS``. + """ policy_map = _maybe_mapping(policy) if policy_map is None: raise ValueError("diffusion_step_policy must be a mapping") - schema_version = policy_map.get("schema_version", _SUPPORTED_SCHEMA_VERSION) + missing = [name for name in _REQUIRED_POLICY_FIELDS if name not in policy_map] + if missing: + raise ValueError(f"diffusion_step_policy missing required fields: {missing}") + + schema_version = policy_map["schema_version"] if schema_version != _SUPPORTED_SCHEMA_VERSION: raise ValueError( f"Unsupported diffusion_step_policy.schema_version={schema_version}; " f"Diffusers supports version {_SUPPORTED_SCHEMA_VERSION}" ) - policy_type = policy_map.get("type") + policy_type = policy_map["type"] if policy_type != _SUPPORTED_POLICY_TYPE: raise ValueError( - f"Unsupported diffusion_step_policy.type={policy_type!r}; " - f"expected {_SUPPORTED_POLICY_TYPE!r}" + f"Unsupported diffusion_step_policy.type={policy_type!r}; expected {_SUPPORTED_POLICY_TYPE!r}" ) - index_space = policy_map.get("index_space", _SUPPORTED_INDEX_SPACE) + index_space = policy_map["index_space"] if index_space != _SUPPORTED_INDEX_SPACE: raise ValueError( - f"Unsupported diffusion_step_policy.index_space={index_space!r}; " - f"expected {_SUPPORTED_INDEX_SPACE!r}" + f"Unsupported diffusion_step_policy.index_space={index_space!r}; expected {_SUPPORTED_INDEX_SPACE!r}" + ) + default_mode = policy_map["default_mode"] + if default_mode != _SUPPORTED_DEFAULT_MODE: + raise ValueError( + f"Unsupported diffusion_step_policy.default_mode={default_mode!r}; expected {_SUPPORTED_DEFAULT_MODE!r}" + ) + scope = policy_map["scope"] + if not isinstance(scope, (list, tuple)) or not all(isinstance(item, str) for item in scope): + raise ValueError("diffusion_step_policy.scope must be a list of strings") + if "transformer" not in scope: + raise ValueError("diffusion_step_policy.scope must include 'transformer'") + unknown_scope = sorted({item for item in scope if item != "transformer"}) + if unknown_scope: + raise ValueError( + f"diffusion_step_policy.scope currently supports only 'transformer', got extra {unknown_scope}" ) if quantization_config: - algo = str( - quantization_config.get("quant_algo") - or quantization_config.get("quant_type") - or "" - ).upper() + algo = str(quantization_config.get("quant_algo") or quantization_config.get("quant_type") or "").upper() if algo and "FP8" not in algo: raise ValueError( "Cosmos3 mixed precision in Diffusers currently supports ModelOpt FP8 " f"checkpoints, got quant_algo/quant_type={algo!r}" ) - first_steps = _window_count(policy_map.get("first_steps"), "first_steps") - last_steps = _window_count(policy_map.get("last_steps"), "last_steps") - overlap = _validated_overlap(str(policy_map.get("overlap", "a16"))) - reasoner = policy_map.get("reasoner", "a16") + first_steps = _window_count(policy_map["first_steps"], "first_steps") + last_steps = _window_count(policy_map["last_steps"], "last_steps") + overlap = _validated_overlap(str(policy_map["overlap"])) + reasoner = policy_map["reasoner"] if reasoner == "a16": reasoner_policy: ReasonerPolicy = "high_precision" elif reasoner == "native": @@ -506,4 +553,3 @@ def _window_count(spec: Any, name: str) -> int: if mode == "native": return 0 return count - diff --git a/tests/pipelines/cosmos/test_cosmos3_mixed_precision.py b/tests/pipelines/cosmos/test_cosmos3_mixed_precision.py index def70d4bc44b..b9b8fbad4c2d 100644 --- a/tests/pipelines/cosmos/test_cosmos3_mixed_precision.py +++ b/tests/pipelines/cosmos/test_cosmos3_mixed_precision.py @@ -1,4 +1,4 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. +# Copyright 2026 The NVIDIA Team and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -157,6 +157,18 @@ def test_reasoner_can_use_base_precision(self): apply_cosmos3_mixed_precision_step(transformer, config, 2, 5) torch.testing.assert_close(reasoner(inputs), expected_w8a8) + def test_dispatches_w8a16_with_fp32_activations(self): + config = Cosmos3MixedPrecisionConfig(format="fp8", first_steps=1, last_steps=0) + transformer = _Transformer() + generation = transformer.layers[0].self_attn.to_q + inputs = torch.tensor([[1.1, -0.7]], dtype=torch.float32) + expected_w8a16 = F.linear( + inputs, + generation.weight.to(inputs.dtype) * generation.weight_quantizer._scale.to(inputs.dtype), + ) + apply_cosmos3_mixed_precision_step(transformer, config, 0, 3) + torch.testing.assert_close(generation(inputs), expected_w8a16) + class Cosmos3CheckpointPolicyTests(unittest.TestCase): def _transformer_with_quant_config(self, quantization_config): @@ -172,7 +184,11 @@ def test_official_policy_enables_first_last_three(self): self.assertEqual(parsed.overlap, "a16") transformer = self._transformer_with_quant_config( - {"quant_method": "modelopt", "quant_algo": "FP8", "runtime": {"diffusion_step_policy": FIRST_LAST_N_FP8_POLICY}} + { + "quant_method": "modelopt", + "quant_algo": "FP8", + "runtime": {"diffusion_step_policy": FIRST_LAST_N_FP8_POLICY}, + } ) config = Cosmos3MixedPrecisionConfig.resolve(transformer) self.assertTrue(config.enabled) @@ -189,14 +205,22 @@ def test_missing_runtime_stays_native_including_four_step_overlap(self): def test_explicit_none_disables_checkpoint_policy(self): transformer = self._transformer_with_quant_config( - {"quant_method": "modelopt", "quant_algo": "FP8", "runtime": {"diffusion_step_policy": FIRST_LAST_N_FP8_POLICY}} + { + "quant_method": "modelopt", + "quant_algo": "FP8", + "runtime": {"diffusion_step_policy": FIRST_LAST_N_FP8_POLICY}, + } ) config = Cosmos3MixedPrecisionConfig.resolve(transformer, mixed_precision_format="none") self.assertFalse(config.enabled) def test_call_site_overrides_checkpoint_counts(self): transformer = self._transformer_with_quant_config( - {"quant_method": "modelopt", "quant_algo": "FP8", "runtime": {"diffusion_step_policy": FIRST_LAST_N_FP8_POLICY}} + { + "quant_method": "modelopt", + "quant_algo": "FP8", + "runtime": {"diffusion_step_policy": FIRST_LAST_N_FP8_POLICY}, + } ) config = Cosmos3MixedPrecisionConfig.resolve( transformer, @@ -215,9 +239,64 @@ def test_force_fp8_without_policy_uses_explicit_schedule(self): self.assertTrue(config.enabled) self.assertEqual([i for i in range(5) if config.use_high_precision(i, 5)], [0, 4]) + def test_policy_on_module_quantization_config_attribute(self): + transformer = _Transformer() + transformer.quantization_config = { + "quant_method": "modelopt", + "quant_algo": "FP8", + "runtime": {"diffusion_step_policy": FIRST_LAST_N_FP8_POLICY}, + } + config = Cosmos3MixedPrecisionConfig.resolve(transformer) + self.assertTrue(config.enabled) + self.assertEqual(config.first_steps, 3) + self.assertEqual(config.last_steps, 3) + + def test_live_null_runtime_overlays_on_disk_policy(self): + class _TransformerWithDiskConfig(_Transformer): + @classmethod + def load_config(cls, name_or_path, local_files_only=True): + return { + "quantization_config": { + "quant_method": "modelopt", + "quant_algo": "FP8", + "runtime": {"diffusion_step_policy": FIRST_LAST_N_FP8_POLICY}, + } + } + + transformer = _TransformerWithDiskConfig() + transformer.quantization_config = { + "quant_method": "modelopt", + "quant_algo": "FP8", + "runtime": None, + } + transformer.config = SimpleNamespace(_name_or_path="/unused/official-fp8") + config = Cosmos3MixedPrecisionConfig.resolve(transformer) + self.assertTrue(config.enabled) + self.assertEqual(config.first_steps, 3) + self.assertEqual(config.last_steps, 3) + + def test_live_null_runtime_without_disk_stays_native(self): + transformer = _Transformer() + transformer.quantization_config = { + "quant_method": "modelopt", + "quant_algo": "FP8", + "runtime": None, + } + transformer.config = SimpleNamespace(_name_or_path=None) + config = Cosmos3MixedPrecisionConfig.resolve(transformer) + self.assertFalse(config.enabled) + def test_malformed_policy_fails_closed(self): with self.assertRaises(ValueError): parse_diffusion_step_policy({"schema_version": 2, "type": "first_last_n"}) + incomplete = dict(FIRST_LAST_N_FP8_POLICY) + del incomplete["overlap"] + with self.assertRaises(ValueError): + parse_diffusion_step_policy(incomplete) + extra_scope = dict(FIRST_LAST_N_FP8_POLICY) + extra_scope["scope"] = ["transformer", "vae"] + with self.assertRaises(ValueError): + parse_diffusion_step_policy(extra_scope) with self.assertRaises(ValueError): Cosmos3MixedPrecisionConfig.resolve( quantization_config={