From b4b907e7b60086c6428145f8dbab687314971dea Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Tue, 18 Aug 2026 10:43:05 -0600 Subject: [PATCH] [Bugfix | Model] Fix SenseNova & Use Well-defined Model Configs (#5877) Signed-off-by: Alex Brooks --- .../test_sensenova_u1_logits_processor.py | 72 ++++++++++ .../sensenova_u1/pipeline_sensenova_u1.py | 113 ++++----------- vllm_omni/diffusion/vllm_config.py | 3 + .../transformers_utils/configs/__init__.py | 3 + .../configs/sensenova_u1.py | 129 ++++++++++++++++++ 5 files changed, 231 insertions(+), 89 deletions(-) create mode 100644 tests/diffusion/models/sensenova_u1/test_sensenova_u1_logits_processor.py create mode 100644 vllm_omni/transformers_utils/configs/sensenova_u1.py diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_logits_processor.py b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_logits_processor.py new file mode 100644 index 000000000..26af9a259 --- /dev/null +++ b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_logits_processor.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression test: SenseNovaU1ForCausalLM instantiation under the diffusion +config shim must not crash on missing ``head_dtype`` and the resulting +LogitsProcessor must have ``head_dtype is None`` (= use model dtype).""" + +import os + +import pytest +import torch +from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config +from vllm.distributed.parallel_state import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, +) + +from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import ( + SenseNovaU1ForCausalLM, +) +from vllm_omni.diffusion.vllm_config import _DiffusionVllmModelConfig +from vllm_omni.transformers_utils.configs.sensenova_u1 import ( + SenseNovaU1Config, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _build_tiny_sensenova_u1(quant_config=None, prefix="model"): + """Build a minimal SenseNovaU1ForCausalLM for unit tests.""" + config = SenseNovaU1Config( + llm_config={ + "hidden_size": 64, + "num_attention_heads": 2, + "num_hidden_layers": 1, + "intermediate_size": 128, + "vocab_size": 32, + "max_position_embeddings": 128, + "max_position_embeddings_hw": 128, + }, + ) + return SenseNovaU1ForCausalLM(config.llm_config, quant_config=quant_config, prefix=prefix) + + +def test_logits_processor_head_dtype_under_diffusion_shim(): + # Initialize the distributed environment + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29543") + init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://") + initialize_model_parallel() + + # Create a vLLM config with a wrapped diffusion config + fake_diff_config = _DiffusionVllmModelConfig( + model="sensenova-u1-test", + dtype=torch.bfloat16, + max_model_len=8192, + original_max_model_len=8192, + ) + vllm_config = VllmConfig( + device_config=DeviceConfig(device="cpu"), + ) + vllm_config.modelconfig = fake_diff_config # type: ignore[assignment] + + # Build a tiny version of the Causal LM component + with set_current_vllm_config(vllm_config): + model = _build_tiny_sensenova_u1() + + # Ensure that we have a set head_dtype attribute. Currently, the head_dtype + # is set to None + head_dtype = model.logits_processor.head_dtype + assert head_dtype is None + cleanup_dist_env_and_memory() diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py index 67e141e6a..3a828fa69 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py +++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py @@ -17,7 +17,6 @@ Key integration points: from __future__ import annotations -import json import math import os from collections.abc import Iterable @@ -28,6 +27,7 @@ import numpy as np import torch import torch.nn as nn import torchvision.transforms as T +from huggingface_hub import snapshot_download from PIL import Image from transformers import AutoTokenizer from vllm.logger import init_logger @@ -40,6 +40,9 @@ from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.transformers_utils.configs.sensenova_u1 import ( + SenseNovaU1Config, +) from .sensenova_u1_transformer import ( SenseNovaU1ForCausalLM, @@ -86,16 +89,12 @@ SYSTEM_MESSAGE_FOR_GEN = ( # --------------------------------------------------------------------------- -# Config parsed from config.json +# Config # --------------------------------------------------------------------------- - - def _resolve_model_path(model_path: str) -> str: """Resolve a HuggingFace model ID or local path to a local directory.""" if os.path.isdir(model_path): return model_path - from huggingface_hub import snapshot_download - return snapshot_download(model_path) @@ -182,73 +181,6 @@ def _load_image_native( return flatten_pv, grid_hw -def _load_sensenova_u1_config(model_path: str): - """Parse config.json and return (top_config, llm_config, vision_config) namespaces.""" - local_path = _resolve_model_path(model_path) - cfg_path = os.path.join(local_path, "config.json") - with open(cfg_path, encoding="utf-8") as f: - raw = json.load(f) - - llm_raw = raw.get("llm_config", {}) - vis_raw = raw.get("vision_config", {}) - - # LLM config needs these extra fields for 3D RoPE - llm_cfg = SimpleNamespace( - hidden_size=llm_raw.get("hidden_size", 4096), - intermediate_size=llm_raw.get("intermediate_size", 11008), - num_hidden_layers=llm_raw.get("num_hidden_layers", 32), - num_attention_heads=llm_raw.get("num_attention_heads", 32), - num_key_value_heads=llm_raw.get("num_key_value_heads", llm_raw.get("num_attention_heads", 32)), - head_dim=llm_raw.get("head_dim", llm_raw.get("hidden_size", 4096) // llm_raw.get("num_attention_heads", 32)), - hidden_act=llm_raw.get("hidden_act", "silu"), - rms_norm_eps=llm_raw.get("rms_norm_eps", 1e-6), - vocab_size=llm_raw.get("vocab_size", 152064), - rope_theta=llm_raw.get("rope_theta", 1000000.0), - rope_theta_hw=llm_raw.get("rope_theta_hw", 10000.0), - max_position_embeddings=llm_raw.get("max_position_embeddings", 40960), - max_position_embeddings_hw=llm_raw.get("max_position_embeddings_hw", 10000), - attention_bias=llm_raw.get("attention_bias", True), - layer_types=llm_raw.get("layer_types", ["full_attention"] * llm_raw.get("num_hidden_layers", 32)), - tie_word_embeddings=llm_raw.get("tie_word_embeddings", False), - ) - - vis_cfg = SimpleNamespace( - num_channels=vis_raw.get("num_channels", 3), - patch_size=vis_raw.get("patch_size", 16), - hidden_size=vis_raw.get("hidden_size", 1024), - llm_hidden_size=vis_raw.get("llm_hidden_size", [llm_cfg.hidden_size]), - downsample_ratio=vis_raw.get("downsample_ratio", [0.5]), - rope_theta_vision=vis_raw.get("rope_theta_vision", 10000.0), - max_position_embeddings_vision=vis_raw.get("max_position_embeddings_vision", 10000), - output_hidden_states=vis_raw.get("output_hidden_states", False), - use_return_dict=vis_raw.get("use_return_dict", True), - ) - - top_cfg = SimpleNamespace( - downsample_ratio=raw.get("downsample_ratio", 0.5), - template=raw.get("template", "neo1_0"), - fm_head_layers=raw.get("fm_head_layers", 2), - fm_head_dim=raw.get("fm_head_dim", 4096), - fm_head_mlp_ratio=raw.get("fm_head_mlp_ratio", 1.0), - use_pixel_head=raw.get("use_pixel_head", False), - noise_scale=raw.get("noise_scale", 1.0), - noise_scale_mode=raw.get("noise_scale_mode", "none"), - noise_scale_base_image_seq_len=raw.get("noise_scale_base_image_seq_len", 256), - noise_scale_max_value=raw.get("noise_scale_max_value", 10.0), - add_noise_scale_embedding=raw.get("add_noise_scale_embedding", False), - time_schedule=raw.get("time_schedule", "standard"), - time_shift_type=raw.get("time_shift_type", "exponential"), - base_shift=raw.get("base_shift", 0.5), - max_shift=raw.get("max_shift", 1.15), - base_image_seq_len=raw.get("base_image_seq_len", 256), - max_image_seq_len=raw.get("max_image_seq_len", 4096), - concat_time_token_num=raw.get("concat_time_token_num", 0), - t_eps=raw.get("t_eps", 0.02), - ) - - return top_cfg, llm_cfg, vis_cfg - - # --------------------------------------------------------------------------- # Vision Embedding — 2D RoPE + Conv patch embed (no transformer blocks) # --------------------------------------------------------------------------- @@ -549,7 +481,9 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli model_path = od_config.model self.local_model_path = _resolve_model_path(model_path) - self.top_cfg, self.llm_cfg, self.vis_cfg = _load_sensenova_u1_config(model_path) + self.model_cfg = SenseNovaU1Config.from_pretrained(self.local_model_path) + self.llm_cfg = self.model_cfg.llm_config + self.vis_cfg = self.model_cfg.vision_config self.tokenizer = AutoTokenizer.from_pretrained(self.local_model_path) self.img_context_token_id = self.tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN) @@ -573,13 +507,13 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli # FM modules llm_hidden_size = self.llm_cfg.hidden_size patch_size = self.vis_cfg.patch_size - merge_size = int(1 / self.top_cfg.downsample_ratio) + merge_size = int(1 / self.model_cfg.downsample_ratio) output_dim = 3 * (patch_size * merge_size) ** 2 vision_model_mot_gen = NEOVisionModel(self.vis_cfg) timestep_embedder = TimestepEmbedder(llm_hidden_size) - if self.top_cfg.use_pixel_head: + if self.model_cfg.use_pixel_head: fm_head = ConvDecoder(llm_hidden_size) else: # The reference implementation uses a fixed 2-layer GELU MLP with a 4096-d @@ -601,13 +535,13 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli } ) - if self.top_cfg.add_noise_scale_embedding: + if self.model_cfg.add_noise_scale_embedding: self.fm_modules["noise_scale_embedder"] = TimestepEmbedder(llm_hidden_size) # Config shortcuts self.patch_size = patch_size self.merge_size = merge_size - self.downsample_ratio = self.top_cfg.downsample_ratio + self.downsample_ratio = self.model_cfg.downsample_ratio # Weight sources for diffusers_loader self.weights_sources = [ @@ -727,6 +661,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli t, z, image_token_num, + t_eps: float, image_size=None, cache_dit_skip=False, **_kw, @@ -748,7 +683,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli ) last_hidden_state = outputs.hidden_states - if self.top_cfg.use_pixel_head: + if self.model_cfg.use_pixel_head: merge_size = self.merge_size token_h = image_size[1] // (self.patch_size * merge_size) token_w = image_size[0] // (self.patch_size * merge_size) @@ -765,7 +700,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli else: x_pred = self.fm_modules["fm_head"](last_hidden_state[:, -image_token_num:].view(B, L, -1)).view(B, L, -1) - v_pred = (x_pred - z) / (1 - t).clamp_min(self.top_cfg.t_eps) + v_pred = (x_pred - z) / (1 - t).clamp_min(t_eps) return v_pred def _apply_time_schedule(self, t, image_seq_len, timestep_shift): @@ -1055,14 +990,14 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli token_w = p.image_size[0] // (self.patch_size * merge_size) grid_hw = torch.tensor([[grid_h, grid_w]] * p.batch_size, device=self.device) - noise_scale = self.top_cfg.noise_scale - if self.top_cfg.noise_scale_mode in ("resolution", "dynamic", "dynamic_sqrt"): - base = float(self.top_cfg.noise_scale_base_image_seq_len) + noise_scale = self.model_cfg.noise_scale + if self.model_cfg.noise_scale_mode in ("resolution", "dynamic", "dynamic_sqrt"): + base = float(self.model_cfg.noise_scale_base_image_seq_len) scale = math.sqrt((grid_h * grid_w) / (merge_size**2) / base) - noise_scale = scale * float(self.top_cfg.noise_scale) - if self.top_cfg.noise_scale_mode == "dynamic_sqrt": + noise_scale = scale * float(self.model_cfg.noise_scale) + if self.model_cfg.noise_scale_mode == "dynamic_sqrt": noise_scale = math.sqrt(noise_scale) - noise_scale = min(noise_scale, self.top_cfg.noise_scale_max_value) + noise_scale = min(noise_scale, self.model_cfg.noise_scale_max_value) generator = torch.Generator(self.device).manual_seed(p.seed) image_prediction = noise_scale * torch.randn( @@ -1101,6 +1036,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli z=z, image_token_num=ns.token_h * ns.token_w, image_size=p.image_size, + t_eps=p.t_eps, ) if cache_dit_skip: kwargs["cache_dit_skip"] = True @@ -1241,7 +1177,6 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli @torch.inference_mode() def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: p = self._parse_request(req) - self.top_cfg.t_eps = p.t_eps input_images = self._extract_input_images(p.first_prompt) modalities = p.first_prompt.get("modalities", []) if isinstance(p.first_prompt, dict) else [] @@ -1438,8 +1373,8 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli ns.token_h * ns.token_w, -1, ) - if self.top_cfg.add_noise_scale_embedding: - ns_tensor = torch.full_like(t_expanded, ns.noise_scale / self.top_cfg.noise_scale_max_value) + if self.model_cfg.add_noise_scale_embedding: + ns_tensor = torch.full_like(t_expanded, ns.noise_scale / self.model_cfg.noise_scale_max_value) ns_emb = self.fm_modules["noise_scale_embedder"](ns_tensor).view( p.batch_size, ns.token_h * ns.token_w, diff --git a/vllm_omni/diffusion/vllm_config.py b/vllm_omni/diffusion/vllm_config.py index 816739de4..b56109de1 100644 --- a/vllm_omni/diffusion/vllm_config.py +++ b/vllm_omni/diffusion/vllm_config.py @@ -81,6 +81,9 @@ class _DiffusionVllmModelConfig: use_mla: bool = False is_moe: bool = False + # Needed for models that bundle things like LogitsProcessors, e.g., SenseNova + head_dtype: torch.dtype | None = None + @property def is_quantized(self) -> bool: return self.quantization is not None diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 2a31b461c..42d6d8d45 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -31,6 +31,7 @@ _CLASS_TO_MODULE: dict[str, str] = { "MingFlashOmniConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", "Qwen3VLMoeVisionConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", "WhisperEncoderConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", + "SenseNovaU1Config": "vllm_omni.transformers_utils.configs.sensenova_u1", } __all__ = [ @@ -55,6 +56,7 @@ __all__ = [ "MingFlashOmniConfig", "Qwen3VLMoeVisionConfig", "WhisperEncoderConfig", + "SenseNovaU1Config", ] @@ -82,5 +84,6 @@ from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2 from vllm_omni.transformers_utils.configs import ming_flash_omni as _ming_flash_omni # noqa: F401, E402 from vllm_omni.transformers_utils.configs import minimax_music3 as _minimax_music3 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import omnivoice as _omnivoice # noqa: F401, E402 +from vllm_omni.transformers_utils.configs import sensenova_u1 as _sensenova_u1 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import voxcpm2 as _voxcpm2 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import voxtral_tts as _voxtral_tts # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/sensenova_u1.py b/vllm_omni/transformers_utils/configs/sensenova_u1.py new file mode 100644 index 000000000..81e8b3406 --- /dev/null +++ b/vllm_omni/transformers_utils/configs/sensenova_u1.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HuggingFace-style configuration classes for SenseNova-U1.""" + +from transformers import AutoConfig, PretrainedConfig, Qwen3Config + +_SENSENOVA_U1_LLM_HIDDEN_DEFAULT = 4096 + + +# Adapted from: https://github.com/OpenSenseNova/SenseNova-U1/blob/main/src/sensenova_u1/models/neo_unify/configuration_neo_chat.py +class SenseNovaU1LLMConfig(Qwen3Config): + """Qwen3-based LLM backbone config with 3D RoPE extensions.""" + + model_type = "sensenova_u1_llm" + + def __init__( + self, + rope_theta: float = 10000.0, + rope_theta_hw: float = 10000.0, + max_position_embeddings_hw: int = 10000, + **kwargs, + ): + self.rope_theta = rope_theta + self.rope_theta_hw = rope_theta_hw + self.max_position_embeddings_hw = max_position_embeddings_hw + super().__init__(**kwargs) + + +# Adapted from https://github.com/OpenSenseNova/SenseNova-U1/blob/main/src/sensenova_u1/models/neo_unify/configuration_neo_vit.py#L10 +class SenseNovaU1VisionConfig(PretrainedConfig): + """Vision embedding config (2D RoPE + conv patch embed, no transformer).""" + + model_type = "sensenova_u1_vision" + + def __init__( + self, + num_channels: int = 3, + patch_size: int = 16, + hidden_size: int = 1024, + llm_hidden_size: list[int] | None = None, + downsample_ratio: list[float] | None = None, + rope_theta_vision: float = 10000.0, + max_position_embeddings_vision: int = 10000, + **kwargs, + ): + self.num_channels = num_channels + self.patch_size = patch_size + self.hidden_size = hidden_size + self.llm_hidden_size = llm_hidden_size if llm_hidden_size is not None else [_SENSENOVA_U1_LLM_HIDDEN_DEFAULT] + self.downsample_ratio = downsample_ratio if downsample_ratio is not None else [0.5] + self.rope_theta_vision = rope_theta_vision + self.max_position_embeddings_vision = max_position_embeddings_vision + super().__init__(**kwargs) + + +# Adapted from https://github.com/OpenSenseNova/SenseNova-U1/blob/main/training/sensenovavl/model/sensenovavl_moe_chat/configuration_sensenovavl_chat.py#L17 +class SenseNovaU1Config(PretrainedConfig): + """Top-level composite config for SenseNova-U1. + + Nests ``llm_config`` and ``vision_config`` sub-configs alongside the + flow-matching / diffusion parameters. When constructed from a dict + (e.g. via ``from_pretrained``), sub-dicts are automatically promoted + to their typed config objects. + """ + + model_type = "sensenova_u1" + + def __init__( + self, + llm_config: dict | SenseNovaU1LLMConfig | None = None, + vision_config: dict | SenseNovaU1VisionConfig | None = None, + downsample_ratio: float = 0.5, + template: str = "neo1_0", + fm_head_layers: int = 2, + fm_head_dim: int = 4096, + fm_head_mlp_ratio: float = 1.0, + use_pixel_head: bool = False, + noise_scale: float = 1.0, + noise_scale_mode: str = "none", + noise_scale_base_image_seq_len: int = 256, + noise_scale_max_value: float = 10.0, + add_noise_scale_embedding: bool = False, + time_schedule: str = "standard", + time_shift_type: str = "exponential", + base_shift: float = 0.5, + max_shift: float = 1.15, + base_image_seq_len: int = 256, + max_image_seq_len: int = 4096, + concat_time_token_num: int = 0, + t_eps: float = 0.02, + **kwargs, + ): + if isinstance(llm_config, dict): + llm_config = SenseNovaU1LLMConfig(**llm_config) + self.llm_config = llm_config or SenseNovaU1LLMConfig() + + if isinstance(vision_config, dict): + vision_config = SenseNovaU1VisionConfig(**vision_config) + self.vision_config = vision_config or SenseNovaU1VisionConfig() + + self.downsample_ratio = downsample_ratio + self.template = template + self.fm_head_layers = fm_head_layers + self.fm_head_dim = fm_head_dim + self.fm_head_mlp_ratio = fm_head_mlp_ratio + self.use_pixel_head = use_pixel_head + self.noise_scale = noise_scale + self.noise_scale_mode = noise_scale_mode + self.noise_scale_base_image_seq_len = noise_scale_base_image_seq_len + self.noise_scale_max_value = noise_scale_max_value + self.add_noise_scale_embedding = add_noise_scale_embedding + self.time_schedule = time_schedule + self.time_shift_type = time_shift_type + self.base_shift = base_shift + self.max_shift = max_shift + self.base_image_seq_len = base_image_seq_len + self.max_image_seq_len = max_image_seq_len + self.concat_time_token_num = concat_time_token_num + self.t_eps = t_eps + super().__init__(**kwargs) + + +AutoConfig.register("sensenova_u1", SenseNovaU1Config) + +__all__ = [ + "SenseNovaU1Config", + "SenseNovaU1LLMConfig", + "SenseNovaU1VisionConfig", +]