diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index ca3c38a91..9f99592fd 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -61,6 +61,7 @@ th { | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | | `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | +| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | | `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | | `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index d01b2973f..81acebe00 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -556,9 +556,21 @@ def main(): elif init_non_diffusion and hasattr(params, "extra_args"): if params.extra_args is None: params.extra_args = {} + params.extra_args.update(diffusion_params.extra_args or {}) if args.seed is not None and hasattr(params, "seed"): params.seed = args.seed + # MammothModa2's AR stage emits one visual token per grid cell, + # one EOL token per row, and one final look-ahead token whose hidden + # state is unavailable. Size the first stage from the prompt metadata + # instead of SamplingParams' default of 16 tokens. + prompt_info = prompt_dict.get("additional_information", {}) + if idx == 0 and prompt_info.get("omni_task") == ["t2i"]: + ar_width = int(prompt_info.get("ar_width", [0])[0]) + ar_height = int(prompt_info.get("ar_height", [0])[0]) + if ar_width > 0 and ar_height > 0: + params.max_tokens = ar_height * (ar_width + 1) + 1 + if not diffusion_replaced and len(sampling_params_list) == 1: sampling_params_list = [diffusion_params] diff --git a/recipes/MammothModa2/MammothModa2-Preview.md b/recipes/MammothModa2/MammothModa2.md similarity index 69% rename from recipes/MammothModa2/MammothModa2-Preview.md rename to recipes/MammothModa2/MammothModa2.md index c94560254..60c391dec 100644 --- a/recipes/MammothModa2/MammothModa2-Preview.md +++ b/recipes/MammothModa2/MammothModa2.md @@ -1,13 +1,12 @@ -# MammothModa2-Preview +# MammothModa2 -> MammothModa2-Preview generation and understanding through shared offline examples +> MammothModa2-Preview and MammothModa2-Dev unified understanding and generation ## Summary - Vendor: ByteDance Research -- Model: `bytedance-research/MammothModa2-Preview` -- Task: Text-to-image generation (AR → DiT two-stage pipeline), plus - image-to-text understanding (AR stage) +- Models: `bytedance-research/MammothModa2-Preview`, `bytedance-research/MammothModa2-Dev` +- Tasks: Preview and Dev text-to-image (AR → DiT); Dev text/image understanding - Mode: Offline inference - Maintainer: Community @@ -29,7 +28,9 @@ flags. Image size uses the standard `--height` / `--width` flags. - Upstream model: [`bytedance-research/MammothModa2-Preview`](https://huggingface.co/bytedance-research/MammothModa2-Preview) -- Related text-to-image example: +- Dev model: + [`bytedance-research/MammothModa2-Dev`](https://huggingface.co/bytedance-research/MammothModa2-Dev) +- Related offline example: [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py) - Related T2T/I2T example: [`examples/offline_inference/x_to_text/x_to_text.py`](../../examples/offline_inference/x_to_text/x_to_text.py) @@ -131,10 +132,24 @@ ls -lh mammoth_t2i.png python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)" ``` -### Image-to-Text (I2T, simplest) +## MammothModa2-Dev unified inference -For the simplest offline image-understanding run, use a short inline Python -snippet from the repository root: +MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses +Qwen2.5-VL. vLLM-Omni selects the matching implementation from the nested +`llm_config.model_type`; no checkpoint edits or `trust_remote_code` flag are +required. + +Text-to-text and image-to-text use the AR-only deploy. Text-to-image loads the +Qwen3 generation experts (`gen_mlp`), extra visual vocabulary and image head, +then sends the generated visual tokens and hidden states to the DiT stage. + +Download the checkpoint: + +```bash +hf download bytedance-research/MammothModa2-Dev --local-dir ./MammothModa2-Dev +``` + +Run image-to-text inference from the repository root: ```python from PIL import Image @@ -142,35 +157,44 @@ from vllm import SamplingParams from vllm.multimodal.image import convert_image_mode from vllm_omni import Omni -def main(): - prompt = ( - "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" - "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" - "Summarize this image.<|im_end|>\n" - "<|im_start|>assistant\n" - ) +prompt = ( + "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" + "Summarize this image.<|im_end|>\n" + "<|im_start|>assistant\n" +) - omni = Omni( - model="./MammothModa2-Preview", - deploy_config="vllm_omni/deploy/mammoth_moda2_ar.yaml", - ) - try: - outputs = list( - omni.generate( - [{ - "prompt": prompt, - "multi_modal_data": {"image": convert_image_mode(Image.open("./image.png"), "RGB")}, - "additional_information": {"omni_task": ["chat"]}, - }], - [SamplingParams(temperature=0.2, top_p=0.9, top_k=-1, max_tokens=512, seed=42)], - ) +omni = Omni( + model="./MammothModa2-Dev", + deploy_config="vllm_omni/deploy/mammoth_moda2_ar.yaml", +) +try: + outputs = list( + omni.generate( + [{ + "prompt": prompt, + "multi_modal_data": { + "image": convert_image_mode(Image.open("./image.png"), "RGB"), + }, + "additional_information": {"omni_task": ["chat"]}, + }], + [SamplingParams( + temperature=0.2, + top_p=0.9, + top_k=-1, + max_tokens=512, + seed=42, + )], ) - finally: - omni.close() + ) +finally: + omni.close() - ro = getattr(outputs[-1], "request_output", outputs[-1]) - print(ro.outputs[0].text.strip()) - -if __name__ == "__main__": - main() +request_output = getattr(outputs[-1], "request_output", outputs[-1]) +print(request_output.outputs[0].text.strip()) ``` + +The Dev checkpoint is approximately 47.55 GiB on disk. In the verified AR-only +run, loaded model weights used approximately 16.97 GiB of GPU memory before KV +and encoder caches. Allow additional GPU memory for those caches and the input +image. diff --git a/recipes/README.md b/recipes/README.md index 13c2320be..94273caa3 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -48,7 +48,7 @@ recipes/ | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB | | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB | | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) | -| [`MammothModa2/MammothModa2-Preview.md`](./MammothModa2/MammothModa2-Preview.md) | Text-to-image with the shared offline image example (AR → DiT) | 1x L40S 48GB / 1x ≥40GB GPU | +| [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom | | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB | | [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) | | [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 | diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index 7c03616bc..47c108235 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -533,15 +533,29 @@ def test_declared_extra_args_apply_to_existing_sampling_params() -> None: @pytest.mark.core_model @pytest.mark.cpu def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> None: - assert get_extra_body_params("MammothModa2DiTPipeline") == frozenset( - { - "text_guidance_scale", - "cfg_range", - "num_inference_steps", - } + for model_class_name in ( + "MammothModa2DiTPipeline", + "MammothModa2ForConditionalGeneration", + "Mammothmoda2Model", + ): + assert get_extra_body_params(model_class_name) == frozenset( + { + "text_guidance_scale", + "cfg_range", + "num_inference_steps", + } + ) + assert get_extra_output_params(model_class_name) == frozenset() + assert should_init_extra_args_for_non_diffusion_stages(model_class_name) is True + + wrapper_prompt = build_text_to_image_prompt( + "MammothModa2ForConditionalGeneration", + prompt="a cat", + negative_prompt=None, + height=256, + width=256, ) - assert get_extra_output_params("MammothModa2DiTPipeline") == frozenset() - assert should_init_extra_args_for_non_diffusion_stages("MammothModa2DiTPipeline") is True + assert wrapper_prompt["additional_information"]["omni_task"] == ["t2i"] @pytest.mark.core_model @@ -568,5 +582,8 @@ def test_mammothmoda2_text_to_image_prompt_builder() -> None: "ar_height": [32], "image_height": [512], "image_width": [768], + "eol_token_id": [152064], + "visual_token_start_id": [152072], + "visual_token_end_id": [168456], }, } diff --git a/tests/unit/mammoth_moda2/test_mammoth_moda2_config.py b/tests/unit/mammoth_moda2/test_mammoth_moda2_config.py index 77b68ca3f..6e35a94f6 100644 --- a/tests/unit/mammoth_moda2/test_mammoth_moda2_config.py +++ b/tests/unit/mammoth_moda2/test_mammoth_moda2_config.py @@ -16,8 +16,12 @@ from vllm_omni.transformers_utils.configs.mammoth_moda2 import ( Mammothmoda2Config, Mammothmoda2Qwen2_5_VLConfig, Mammothmoda2Qwen2_5_VLTextConfig, + Mammothmoda2Qwen3VLConfig, + Mammothmoda2Qwen3VLTextConfig, ) +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + @pytest.mark.cpu class TestMammothmoda2Config: @@ -89,3 +93,30 @@ class TestMammothmoda2Qwen2_5_VLConfig: assert hasattr(text_config, "vocab_size") assert hasattr(text_config, "image_token_id") assert hasattr(text_config, "video_token_id") + + +@pytest.mark.cpu +class TestMammothmoda2Qwen3VLConfig: + def test_dev_config_uses_qwen3_vl_subconfigs(self): + config = Mammothmoda2Config( + llm_config={ + "model_type": "mammothmoda2_qwen3_vl", + "text_config": { + "model_type": "mammothmoda2_qwen3_vl_text", + "vocab_size": 151936, + "gen_vocab_size": 32800, + "gen_vocab_start_index": 152064, + }, + "vision_config": { + "model_type": "mammothmoda2_qwen3_vl_vision", + "deepstack_visual_indexes": [8, 16, 24], + }, + } + ) + + assert isinstance(config.llm_config, Mammothmoda2Qwen3VLConfig) + assert isinstance(config.get_text_config(), Mammothmoda2Qwen3VLTextConfig) + # vLLM validates sampling against the complete base + visual vocabulary. + assert config.get_text_config().vocab_size == 152064 + 32800 + assert config.get_text_config().gen_vocab_start_index == 152064 + assert config.llm_config.vision_config.deepstack_visual_indexes == [8, 16, 24] diff --git a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py index ed42e24ca..000abe051 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py +++ b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py @@ -189,6 +189,7 @@ class SimpleQFormerImageRefiner(nn.Module): def __init__( self, hidden_size: int, + output_hidden_size: int | None = None, num_queries: int = 128, num_layers: int = 2, num_heads: int | None = None, @@ -196,6 +197,8 @@ class SimpleQFormerImageRefiner(nn.Module): norm_eps: float = 1e-5, ) -> None: super().__init__() + input_hidden_size = hidden_size + hidden_size = output_hidden_size or hidden_size self.hidden_size = hidden_size self.num_queries = num_queries # ensure num_heads divides hidden_size @@ -203,8 +206,8 @@ class SimpleQFormerImageRefiner(nn.Module): num_heads = max(1, hidden_size // 128) self.num_heads = self._choose_valid_num_heads(hidden_size, num_heads) self.input_proj = nn.Sequential( - Qwen2RMSNorm(hidden_size, eps=norm_eps), - nn.Linear(hidden_size, hidden_size, bias=True), + Qwen2RMSNorm(input_hidden_size, eps=norm_eps), + nn.Linear(input_hidden_size, hidden_size, bias=True), ) # Learnable query embeddings @@ -662,11 +665,32 @@ class Transformer2DModel(ModelMixin, ConfigMixin): batch_size: int, height: int, width: int, + ar_image_hidden_states: torch.Tensor | None = None, + ar_image_attention_mask: torch.Tensor | None = None, ): device = hidden_states.device p = self.config.patch_size temb, text_hidden_states = self.time_caption_embed(timestep, text_hidden_states, hidden_states.dtype) + image_embedder = getattr(self.time_caption_embed, "image_embedder", None) + if image_embedder is not None and ar_image_hidden_states is not None: + if ar_image_attention_mask is None: + ar_image_attention_mask = torch.ones( + ar_image_hidden_states.shape[:2], + dtype=torch.bool, + device=ar_image_hidden_states.device, + ) + image_hidden_states = image_embedder( + ar_image_hidden_states, + ~ar_image_attention_mask.bool(), + ) + image_attention_mask = torch.ones( + image_hidden_states.shape[:2], + dtype=torch.bool, + device=image_hidden_states.device, + ) + text_hidden_states = torch.cat([text_hidden_states, image_hidden_states], dim=1) + text_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1) img_tokens = rearrange(hidden_states, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=p, p2=p) img_tokens = self.x_embedder(img_tokens) @@ -699,6 +723,7 @@ class Transformer2DModel(ModelMixin, ConfigMixin): return ( temb, text_hidden_states, + text_attention_mask, img_tokens, img_mask, img_len, @@ -740,6 +765,8 @@ class Transformer2DModel(ModelMixin, ConfigMixin): freqs_cis: torch.Tensor, text_attention_mask: torch.Tensor, ref_image_hidden_states: list[list[torch.Tensor]] | None = None, + ar_image_hidden_states: torch.Tensor | None = None, + ar_image_attention_mask: torch.Tensor | None = None, return_dict: bool = False, ) -> torch.Tensor: batch_size, height, width = self._validate_inputs( @@ -749,6 +776,7 @@ class Transformer2DModel(ModelMixin, ConfigMixin): ( temb, text_hidden_states, + text_attention_mask, img_tokens, img_mask, img_len, @@ -766,6 +794,8 @@ class Transformer2DModel(ModelMixin, ConfigMixin): batch_size, height, width, + ar_image_hidden_states, + ar_image_attention_mask, ) text_hidden_states, img_tokens = self._apply_refiners( diff --git a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py index 614c8268c..5298d5514 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py +++ b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py @@ -42,6 +42,7 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ "llm_model.": None, + "gen_tokenizer.": None, } ) @@ -77,11 +78,23 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): ) self._reinit_caption_embedder(llm_hidden_size) - # Optional: image condition refiner (Q-Former) - if self.config.gen_image_condition_refiner_config is not None: + # Optional image condition Q-Former. Preview stores it as a standalone + # module; Dev stores it under the DiT timestep/caption embedder. + llm_model_type = getattr(self.config.llm_config, "model_type", "") + refiner_config = self.config.gen_image_condition_refiner_config + if refiner_config is not None and llm_model_type == "mammothmoda2_qwen3_vl": + dit_hidden_size = int(self.gen_transformer.hidden_size) + self.gen_transformer.time_caption_embed.image_embedder = SimpleQFormerImageRefiner( + hidden_size=llm_hidden_size, + output_hidden_size=dit_hidden_size, + num_heads=max(1, dit_hidden_size // 128), + **refiner_config, + ) + self.gen_image_condition_refiner = None + elif refiner_config is not None: self.gen_image_condition_refiner = SimpleQFormerImageRefiner( hidden_size=llm_hidden_size, - **self.config.gen_image_condition_refiner_config, + **refiner_config, ) else: self.gen_image_condition_refiner = None @@ -237,6 +250,13 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): text_cond = _ensure_2d(text_cond, "text_prompt_embeds") image_cond = _ensure_2d(image_cond, "image_prompt_embeds") + if image_cond.shape[0] == 0: + answer_token_ids = info.get("full_token_ids", [])[int(info.get("answer_start_index", [0])[0]) :] + raise ValueError( + "MammothModa2 AR stage produced no visual-token hidden states; " + "the DiT stage requires at least one generated visual token. " + f"Generated token ids: {answer_token_ids[:32]}" + ) text_cond = text_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous() image_cond = image_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous() @@ -263,8 +283,17 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): device=image_embeds.device, ) - prompt_embeds = torch.cat([text_embeds, image_embeds], dim=1) - prompt_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1) + nested_image_embedder = getattr(self.gen_transformer.time_caption_embed, "image_embedder", None) + if nested_image_embedder is None: + prompt_embeds = torch.cat([text_embeds, image_embeds], dim=1) + prompt_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1) + ar_image_embeds = None + ar_image_attention_mask = None + else: + prompt_embeds = text_embeds + prompt_attention_mask = text_attention_mask + ar_image_embeds = image_embeds + ar_image_attention_mask = image_attention_mask # Prepare negative prompt (for CFG). If none provided, fall back to unconditional. negative_prompt_embeds = None @@ -337,6 +366,8 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): text_hidden_states=prompt_embeds, text_attention_mask=prompt_attention_mask, ref_image_hidden_states=None, + ar_image_hidden_states=ar_image_embeds, + ar_image_attention_mask=ar_image_attention_mask, freqs_cis=self.gen_freqs_cis, ) guidance_scale = text_guidance_scale if cfg_range[0] <= i / total_steps <= cfg_range[1] else 1.0 diff --git a/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py b/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py index 2fe616cfd..47c2b5e37 100644 --- a/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py +++ b/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py @@ -6,7 +6,7 @@ from typing import Any import torch from torch import nn -from transformers import Qwen2Config +from transformers import Qwen2Config, Qwen3VLProcessor from transformers.models.qwen2_5_vl.processing_qwen2_5_vl import Qwen2_5_VLProcessor from vllm.config import CacheConfig, VllmConfig from vllm.distributed import get_pp_group @@ -30,7 +30,15 @@ from vllm.model_executor.models.qwen2_5_vl import ( Qwen2_5_VLProcessingInfo, ) from vllm.model_executor.models.qwen2_vl import Qwen2VLMultiModalDataParser +from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer, Qwen3MLP +from vllm.model_executor.models.qwen3_vl import ( + Qwen3VLDummyInputsBuilder, + Qwen3VLForConditionalGeneration, + Qwen3VLMultiModalProcessor, + Qwen3VLProcessingInfo, +) from vllm.model_executor.models.utils import ( + AutoWeightsLoader, PPMissingLayer, WeightsMapper, init_vllm_registered_model, @@ -51,6 +59,18 @@ from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config +def _runtime_meta(runtime_info: dict[str, Any]) -> dict[str, Any]: + """Accept both legacy flat metadata and the canonical nested payload.""" + meta = runtime_info.get("meta") + return meta if isinstance(meta, dict) else runtime_info + + +def _first_int(value: Any) -> int: + if isinstance(value, (list, tuple)): + value = value[0] + return int(value) + + def moe_enable(moe_type, layer_type, layer_idx) -> bool: """Determine if MoE should be enabled for a specific layer type and index. @@ -291,6 +311,7 @@ class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP): self.gen_vocab_size = int(getattr(config, "gen_vocab_size", 0) or 0) self.base_vocab_size = int(self.gen_vocab_start_index) if self.extra_gen_vocab else int(config.vocab_size) + self.org_base_vocab_size = int(getattr(config, "base_vocab_size", self.base_vocab_size)) # The configuration level (hf_text_config.vocab_size) has been extended to base+gen # by the upstream config class. Use config.vocab_size as the total vocab size. self.total_vocab_size = int(getattr(config, "vocab_size", self.base_vocab_size)) @@ -299,6 +320,7 @@ class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP): self.embed_tokens = VocabParallelEmbedding( self.base_vocab_size, config.hidden_size, + org_num_embeddings=self.org_base_vocab_size, quant_config=quant_config, prefix=f"{prefix}.embed_tokens", ) @@ -322,6 +344,7 @@ class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP): self.lm_head = ParallelLMHead( self.base_vocab_size, config.hidden_size, + org_num_embeddings=self.org_base_vocab_size, quant_config=quant_config, prefix=f"{prefix}.lm_head", ) @@ -425,6 +448,7 @@ class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP): positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, + deepstack_input_embeds: IntermediateTensors | None = None, ) -> torch.Tensor | IntermediateTensors: if get_pp_group().is_first_rank: if inputs_embeds is not None: @@ -446,8 +470,17 @@ class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP): residual = intermediate_tensors["residual"] gen_token_mask = intermediate_tensors.tensors.get("gen_token_mask") - for idx, layer in enumerate(islice(self.layers, self.start_layer, self.end_layer)): + for idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer), + start=self.start_layer, + ): hidden_states, residual = layer(positions, hidden_states, residual, gen_token_mask) + # The vision encoder's deepstack_visual_indexes select which vision + # layers produce these features. The resulting tensors are keyed + # 0..N-1 and, as in Qwen3-Omni, are injected into the first N + # language-model layers. + if deepstack_input_embeds is not None and idx < len(deepstack_input_embeds): + hidden_states = hidden_states + deepstack_input_embeds[f"deepstack_input_embeds_{idx}"] if not get_pp_group().is_last_rank: tensors = {"hidden_states": hidden_states, "residual": residual} @@ -532,6 +565,65 @@ class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP): return loaded_params +class Mammoth3DecoderLayer(Qwen3DecoderLayer): + """Qwen3 decoder layer with MammothModa2's generation MLP expert.""" + + def __init__( + self, + config, + layer_idx: int, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__(config, cache_config, quant_config, prefix) + self.moe_enable = moe_enable(config.moe_type, "ffn", layer_idx) + self.gen_mlp = ( + Qwen3MLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.gen_mlp", + ) + if self.moe_enable + else None + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + gen_token_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = moe_forward( + hidden_states, + self.mlp, + self.gen_mlp, + gen_token_mask if self.moe_enable else None, + ) + return hidden_states, residual + + +class MammothModa2Qwen3ForCausalLM(MammothModa2Qwen2ForCausalLM): + """MammothModa2 generation-aware language model using Qwen3 attention.""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__( + vllm_config=vllm_config, + prefix=prefix, + decoder_layer_type=Mammoth3DecoderLayer, + ) + + @MULTIMODAL_REGISTRY.register_processor( MammothModa2ARMultiModalProcessor, info=MammothModa2ARProcessingInfo, @@ -604,17 +696,17 @@ class MammothModa2ARForConditionalGeneration(Qwen2_5_VLForConditionalGeneration) num_reqs = int(logits.shape[0]) for i in range(num_reqs): runtime_info = runtime_infos[i] if isinstance(runtime_infos[i], dict) else {} - meta = runtime_info.get("meta", {}) + meta = _runtime_meta(runtime_info) omni_task = meta.get("omni_task") if not isinstance(omni_task, list) or not omni_task or omni_task[0] != "t2i": # Text/understanding/chat: forbid sampling from the extra gen vocab. logits[i, self.language_model.base_vocab_size :] = neg_inf continue - ar_width = meta["ar_width"][0] - eol_token_id = meta["eol_token_id"][0] - visual_start = meta["visual_token_start_id"][0] - visual_end = meta["visual_token_end_id"][0] + ar_width = _first_int(meta["ar_width"]) + eol_token_id = _first_int(meta["eol_token_id"]) + visual_start = _first_int(meta["visual_token_start_id"]) + visual_end = _first_int(meta["visual_token_end_id"]) generated_len = runtime_info["generated_len"] row = logits[i] @@ -682,6 +774,100 @@ class MammothModa2ARForConditionalGeneration(Qwen2_5_VLForConditionalGeneration) return logits +class MammothModa2Qwen3ProcessingInfo(Qwen3VLProcessingInfo): + """Qwen3-VL processing using MammothModa2's nested AR configuration.""" + + def get_hf_config(self): + mammoth_cfg: Mammothmoda2Config = self.ctx.get_hf_config(Mammothmoda2Config) + return mammoth_cfg.llm_config + + def get_hf_processor(self, **kwargs: object) -> Qwen3VLProcessor: + # The MammothModa2 repository does not ship tokenizer_config.json, so + # AutoProcessor otherwise resolves only the image processor and leaves + # Qwen3-VL's image placeholder undefined. + return self.ctx.get_hf_processor( + Qwen3VLProcessor, + use_fast=kwargs.pop("use_fast", True), + **kwargs, + ) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"image": None} + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen3VLMultiModalProcessor, + info=MammothModa2Qwen3ProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) +class MammothModa2Qwen3ARForConditionalGeneration(Qwen3VLForConditionalGeneration): + """MammothModa2-Dev Qwen3-VL stage for understanding and generation.""" + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "llm_model.model.visual.": "visual.", + "llm_model.lm_head.": "language_model.lm_head.", + "llm_model.model.language_model.gen_embed_tokens.": "language_model.gen_embed_tokens.", + "llm_model.model.language_model.": "language_model.", + "llm_model.gen_head.": "language_model.gen_head.", + "gen_image_condition_refiner.": None, + "gen_transformer.": None, + "gen_vae.": None, + "gen_tokenizer.": None, + } + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + mammoth_cfg = vllm_config.model_config.hf_config + ar_hf_config = getattr(mammoth_cfg, "llm_config", mammoth_cfg) + ar_vllm_config = vllm_config.with_hf_config( + ar_hf_config, + architectures=vllm_config.model_config.architectures, + ) + super().__init__(vllm_config=ar_vllm_config, prefix=prefix) + text_config = ar_hf_config.text_config + self.language_model = MammothModa2Qwen3ForCausalLM( + vllm_config=ar_vllm_config.with_hf_config(text_config), + prefix=maybe_prefix(prefix, "language_model"), + ) + self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors + self._last_runtime_additional_information: list[dict[str, Any]] | None = None + + def forward(self, *args, **kwargs): + runtime_infos = kwargs.get("runtime_additional_information") + self._last_runtime_additional_information = runtime_infos if isinstance(runtime_infos, list) else None + return super().forward(*args, **kwargs) + + def compute_logits(self, hidden_states: torch.Tensor): + logits = super().compute_logits(hidden_states) + if not isinstance(logits, torch.Tensor) or self._last_runtime_additional_information is None: + return logits + neg_inf = -float("inf") + for i, runtime_info in enumerate(self._last_runtime_additional_information): + meta = _runtime_meta(runtime_info) if isinstance(runtime_info, dict) else {} + if (meta.get("omni_task") or [None])[0] != "t2i": + logits[i, self.language_model.base_vocab_size :] = neg_inf + continue + ar_width = _first_int(meta["ar_width"]) + eol_token_id = _first_int(meta["eol_token_id"]) + visual_start = _first_int(meta["visual_token_start_id"]) + visual_end = _first_int(meta["visual_token_end_id"]) + row = logits[i] + if int(runtime_info["generated_len"]) % (ar_width + 1) == ar_width: + eol_logit = row[eol_token_id].clone() + row.fill_(neg_inf) + row[eol_token_id] = eol_logit + else: + row[:visual_start] = neg_inf + row[visual_end + 1 :] = neg_inf + row[eol_token_id] = neg_inf + return logits + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + @MULTIMODAL_REGISTRY.register_processor( MammothModa2ARMultiModalProcessor, info=MammothModa2ARProcessingInfo, @@ -716,7 +902,11 @@ class MammothModa2ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor vllm_config=vllm_config, prefix=maybe_prefix(prefix, "ar"), hf_config=cfg.llm_config if hasattr(cfg, "llm_config") else cfg.text_config, - architectures=["MammothModa2ARForConditionalGeneration"], + architectures=[ + "MammothModa2Qwen3ARForConditionalGeneration" + if getattr(cfg.llm_config, "model_type", "") == "mammothmoda2_qwen3_vl" + else "MammothModa2ARForConditionalGeneration" + ], ) self.model = self.ar elif self.model_stage == "dit": @@ -750,6 +940,26 @@ class MammothModa2ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor return self.model.get_language_model() return self.model + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + if isinstance(self.model, MammothModa2Qwen3ARForConditionalGeneration): + return self.model.embed_input_ids( + input_ids, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + return SupportsMultiModal.embed_input_ids( + self, + input_ids, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + def get_multimodal_embeddings(self, **kwargs: object): # Backward compatibility: route through embed_multimodal. return self.embed_multimodal(**kwargs) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 6f0302499..2c3a4c1d7 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -99,6 +99,16 @@ _OMNI_MODELS = { "mammoth_moda2", "MammothModa2ARForConditionalGeneration", ), + "MammothModa2Qwen3ARForConditionalGeneration": ( + "mammoth_moda2", + "mammoth_moda2", + "MammothModa2Qwen3ARForConditionalGeneration", + ), + "MammothModa2Qwen3ForCausalLM": ( + "mammoth_moda2", + "mammoth_moda2", + "MammothModa2Qwen3ForCausalLM", + ), "MammothModa2DiTPipeline": ( "mammoth_moda2", "pipeline_mammothmoda2_dit", diff --git a/vllm_omni/model_extras/mammothmodal2_preview.py b/vllm_omni/model_extras/mammothmodal2_preview.py index ddce493f6..96f3fb939 100644 --- a/vllm_omni/model_extras/mammothmodal2_preview.py +++ b/vllm_omni/model_extras/mammothmodal2_preview.py @@ -9,6 +9,11 @@ from typing import Any # bespoke MammothModa2 example used to convert image dimensions -> AR grid dims. _PATCH_SIZE = 16 +# Shared by Preview and Dev (see ``t2i_generation_config.json``). +_EOL_TOKEN_ID = 152064 +_VISUAL_TOKEN_START_ID = 152072 +_VISUAL_TOKEN_END_ID = 168456 + _AR_SYSTEM_PROMPT = "You are a helpful image generator." @@ -26,9 +31,8 @@ def build_text_to_image_prompt( Model-specific sampling knobs (``text_guidance_scale``, ``cfg_range``, ``num_inference_steps``) flow separately via ``extra_body`` -> ``extra_args``. - Config-derived token ids (``eol_token_id``, ``visual_token_start_id``, - ``visual_token_end_id``, ``visual_ids``) are sourced inside the AR stage - rather than passed by the caller -- see the ar2dit stage input processor. + The structural token ids are included in ``additional_information`` so the + AR sampler can apply the same constraints for both Preview and Dev. MammothModa2 t2i uses classifier-free guidance via ``text_guidance_scale`` and has no explicit negative-prompt path, so ``negative_prompt`` is accepted @@ -52,6 +56,9 @@ def build_text_to_image_prompt( "ar_height": [ar_height], "image_height": [h], "image_width": [w], + "eol_token_id": [_EOL_TOKEN_ID], + "visual_token_start_id": [_VISUAL_TOKEN_START_ID], + "visual_token_end_id": [_VISUAL_TOKEN_END_ID], }, } diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 44051c081..edf28101c 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -243,6 +243,11 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { }, } +# Multi-stage discovery reports the top-level wrapper rather than its DiT +# submodule, so both names must resolve to the same request adapters. +_EXTRA_SPECS["MammothModa2ForConditionalGeneration"] = _EXTRA_SPECS["MammothModa2DiTPipeline"] +_EXTRA_SPECS["Mammothmoda2Model"] = _EXTRA_SPECS["MammothModa2DiTPipeline"] + def _get_spec(model_class_name: str | None) -> dict[str, Any] | None: if not model_class_name: diff --git a/vllm_omni/transformers_utils/configs/mammoth_moda2.py b/vllm_omni/transformers_utils/configs/mammoth_moda2.py index 1c6352264..d0848d1ca 100644 --- a/vllm_omni/transformers_utils/configs/mammoth_moda2.py +++ b/vllm_omni/transformers_utils/configs/mammoth_moda2.py @@ -8,6 +8,11 @@ from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import ( Qwen2_5_VLTextConfig, Qwen2_5_VLVisionConfig, ) +from transformers.models.qwen3_vl.configuration_qwen3_vl import ( + Qwen3VLConfig, + Qwen3VLTextConfig, + Qwen3VLVisionConfig, +) from vllm_omni.tokenizers.mammoth_moda2_tokenizer import MammothUTokenizer @@ -16,6 +21,9 @@ __all__ = [ "Mammothmoda2Qwen2_5_VLConfig", "Mammothmoda2Qwen2_5_VLTextConfig", "Mammothmoda2Qwen2_5_VLVisionConfig", + "Mammothmoda2Qwen3VLConfig", + "Mammothmoda2Qwen3VLTextConfig", + "Mammothmoda2Qwen3VLVisionConfig", ] @@ -209,6 +217,77 @@ class Mammothmoda2Qwen2_5_VLConfig(Qwen2_5_VLConfig): self.tokenizer_class = "MammothUTokenizer" +class Mammothmoda2Qwen3VLVisionConfig(Qwen3VLVisionConfig): + model_type = "mammothmoda2_qwen3_vl_vision" + base_config_key = "vision_config" + + +class Mammothmoda2Qwen3VLTextConfig(Qwen3VLTextConfig): + model_type = "mammothmoda2_qwen3_vl_text" + base_config_key = "text_config" + + def __init__( + self, + extra_gen_vocab: bool = True, + gen_vocab_size: int = 32800, + gen_vocab_start_index: int | None = None, + moe_type: str = "ffn", + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.base_vocab_size = int(self.vocab_size) + self.extra_gen_vocab = extra_gen_vocab + self.gen_vocab_size = gen_vocab_size + self.gen_vocab_start_index = self.vocab_size if gen_vocab_start_index is None else gen_vocab_start_index + self.moe_type = moe_type + if self.extra_gen_vocab: + self.vocab_size = int(self.gen_vocab_start_index) + int(self.gen_vocab_size) + + +class Mammothmoda2Qwen3VLConfig(Qwen3VLConfig): + """MammothModa2-Dev AR config backed by Qwen3-VL.""" + + model_type = "mammothmoda2_qwen3_vl" + sub_configs = { + "vision_config": Mammothmoda2Qwen3VLVisionConfig, + "text_config": Mammothmoda2Qwen3VLTextConfig, + } + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + text_config: dict | PretrainedConfig | None = None, + vision_config: dict | PretrainedConfig | None = None, + extra_gen_vocab: bool = True, + gen_vocab_size: int = 32800, + gen_vocab_start_index: int | None = None, + moe_type: str = "ffn", + **kwargs, + ) -> None: + if isinstance(vision_config, dict): + vision_config = self.sub_configs["vision_config"](**vision_config) + elif vision_config is None: + vision_config = self.sub_configs["vision_config"]() + extras = { + "extra_gen_vocab": extra_gen_vocab, + "gen_vocab_size": gen_vocab_size, + "gen_vocab_start_index": gen_vocab_start_index, + "moe_type": moe_type, + } + if isinstance(text_config, dict): + for key, value in extras.items(): + text_config.setdefault(key, value) + text_config = self.sub_configs["text_config"](**text_config) + elif text_config is None: + text_config = self.sub_configs["text_config"](**extras) + super().__init__(text_config=text_config, vision_config=vision_config, **kwargs) + self.extra_gen_vocab = getattr(self.text_config, "extra_gen_vocab", extra_gen_vocab) + self.gen_vocab_size = getattr(self.text_config, "gen_vocab_size", gen_vocab_size) + self.gen_vocab_start_index = getattr(self.text_config, "gen_vocab_start_index", gen_vocab_start_index) + self.moe_type = getattr(self.text_config, "moe_type", moe_type) + self.tokenizer_class = "MammothUTokenizer" + + class Mammothmoda2Config(PretrainedConfig): """Top-level MammothModa2 composition configuration""" @@ -289,6 +368,10 @@ AutoConfig.register(Mammothmoda2Config.model_type, Mammothmoda2Config) AutoConfig.register(Mammothmoda2Qwen2_5_VLConfig.model_type, Mammothmoda2Qwen2_5_VLConfig) AutoConfig.register(Mammothmoda2Qwen2_5_VLTextConfig.model_type, Mammothmoda2Qwen2_5_VLTextConfig) AutoConfig.register(Mammothmoda2Qwen2_5_VLVisionConfig.model_type, Mammothmoda2Qwen2_5_VLVisionConfig) +AutoConfig.register(Mammothmoda2Qwen3VLConfig.model_type, Mammothmoda2Qwen3VLConfig) +AutoConfig.register(Mammothmoda2Qwen3VLTextConfig.model_type, Mammothmoda2Qwen3VLTextConfig) +AutoConfig.register(Mammothmoda2Qwen3VLVisionConfig.model_type, Mammothmoda2Qwen3VLVisionConfig) # Register tokenizer_type -> the configs & AutoTokenizer AutoTokenizer.register(config_class=Mammothmoda2Config, slow_tokenizer_class=MammothUTokenizer) AutoTokenizer.register(config_class=Mammothmoda2Qwen2_5_VLConfig, slow_tokenizer_class=MammothUTokenizer) +AutoTokenizer.register(config_class=Mammothmoda2Qwen3VLConfig, slow_tokenizer_class=MammothUTokenizer)