[Model] Add Bailing V3 VL support (#55921)
Signed-off-by: zexplorerhj <zhjoneson@163.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
zexplorerhj
co-authored by
Codex
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Codex
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent
6ff479e1f7
commit
fe26c705f7
@@ -538,6 +538,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `AriaForConditionalGeneration` | Aria | T + I<sup>+</sup> | `rhymes-ai/Aria` | | |
|
||||
| `AudioFlamingo3ForConditionalGeneration` | AudioFlamingo3 | T + A | `nvidia/audio-flamingo-3-hf`, `nvidia/music-flamingo-hf` | ✅︎ | ✅︎ |
|
||||
| `BagelForConditionalGeneration` | BAGEL | T + I<sup>+</sup> | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ |
|
||||
| `BailingMoeV3VLForConditionalGeneration` | Ling 3.0 Flash VL | T + I<sup>E+</sup> | `inclusionAI/Ling-3.0-flash-VL` | | ✅︎ |
|
||||
| `BeeForConditionalGeneration` | Bee-8B | T + I<sup>E+</sup> | `Open-Bee/Bee-8B-RL`, `Open-Bee/Bee-8B-SFT` | | ✅︎ |
|
||||
| `Blip2ForConditionalGeneration` | BLIP-2 | T + I<sup>E</sup> | `Salesforce/blip2-opt-2.7b`, `Salesforce/blip2-opt-6.7b`, etc. | ✅︎ | ✅︎ |
|
||||
| `Cohere2VisionForConditionalGeneration` | Command A Vision, Command-A+ | T + I<sup>+</sup> | `CohereLabs/command-a-vision-07-2025`, `CohereLabs/command-a-plus-05-2026`, etc. | | ✅︎ |
|
||||
|
||||
@@ -0,0 +1,176 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import copy
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
|
||||
from vllm.model_executor.layers.rotary_embedding import get_rope
|
||||
from vllm.model_executor.layers.rotary_embedding.bailing_mrope import (
|
||||
BailingMRotaryEmbedding,
|
||||
)
|
||||
from vllm.model_executor.models.bailing_moe_v3 import _build_mla_rotary_embedding
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
def test_bailing_mrope_selects_alternating_spatial_then_temporal():
|
||||
with set_current_vllm_config(VllmConfig(device_config=DeviceConfig("cpu"))):
|
||||
rotary_emb = BailingMRotaryEmbedding(
|
||||
head_size=12,
|
||||
rotary_dim=12,
|
||||
max_position_embeddings=8,
|
||||
base=10_000,
|
||||
is_neox_style=False,
|
||||
dtype=torch.float32,
|
||||
mrope_section=[2, 2, 2],
|
||||
)
|
||||
|
||||
positions = torch.tensor([[1, 2], [3, 4], [5, 6]])
|
||||
frequencies = torch.arange(6).expand(8, -1)
|
||||
position_values = torch.arange(8).unsqueeze(1) * 10
|
||||
cos = position_values + frequencies
|
||||
sin = cos + 100
|
||||
cos_sin_cache = torch.cat((cos, sin), dim=-1)
|
||||
|
||||
actual_cos, actual_sin = rotary_emb.select_cos_sin(
|
||||
positions,
|
||||
cos_sin_cache,
|
||||
)
|
||||
|
||||
expected_cos = torch.tensor(
|
||||
[
|
||||
[30, 51, 32, 53, 14, 15],
|
||||
[40, 61, 42, 63, 24, 25],
|
||||
]
|
||||
)
|
||||
torch.testing.assert_close(actual_cos, expected_cos)
|
||||
torch.testing.assert_close(actual_sin, expected_cos + 100)
|
||||
|
||||
|
||||
def test_bailing_factory_preserves_mla_width_and_config():
|
||||
config = SimpleNamespace(
|
||||
max_position_embeddings=128,
|
||||
rope_parameters={
|
||||
"rope_type": "default",
|
||||
"rope_theta": 6_000_000,
|
||||
"partial_rotary_factor": 0.5,
|
||||
"mrope_section": [8, 12, 12],
|
||||
},
|
||||
rope_scaling={"partial_rotary_factor": 0.5},
|
||||
)
|
||||
original = copy.deepcopy(vars(config))
|
||||
with set_current_vllm_config(VllmConfig(device_config=DeviceConfig("cpu"))):
|
||||
rope = _build_mla_rotary_embedding(config, head_size=64)
|
||||
assert rope is _build_mla_rotary_embedding(config, head_size=64)
|
||||
generic = get_rope(
|
||||
head_size=64,
|
||||
max_position=128,
|
||||
is_neox_style=False,
|
||||
rope_parameters={
|
||||
"rope_type": "default",
|
||||
"rope_theta": 6_000_000,
|
||||
"mrope_section": [8, 12, 12],
|
||||
},
|
||||
)
|
||||
assert isinstance(rope, BailingMRotaryEmbedding)
|
||||
assert rope is not generic
|
||||
assert rope.rotary_dim == 64
|
||||
assert vars(config) == original
|
||||
|
||||
|
||||
def test_bailing_factory_rejects_unsupported_scaling():
|
||||
config = SimpleNamespace(
|
||||
rope_parameters={
|
||||
"rope_type": "yarn",
|
||||
"factor": 2.0,
|
||||
"mrope_section": [8, 12, 12],
|
||||
}
|
||||
)
|
||||
with pytest.raises(ValueError, match="only supports rope_type='default'"):
|
||||
_build_mla_rotary_embedding(config, head_size=64)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="Skipping CUDA/ROCm-only test.",
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize("position_kind", ["text", "multimodal"])
|
||||
@torch.inference_mode()
|
||||
def test_bailing_mrope_cuda_matches_native_for_strided_mla_views(
|
||||
dtype: torch.dtype,
|
||||
position_kind: str,
|
||||
):
|
||||
device = torch.device("cuda")
|
||||
num_tokens = 17
|
||||
num_query_heads = 8
|
||||
with set_current_vllm_config(VllmConfig()):
|
||||
rotary_emb = BailingMRotaryEmbedding(
|
||||
head_size=64,
|
||||
rotary_dim=64,
|
||||
max_position_embeddings=64,
|
||||
base=6_000_000,
|
||||
is_neox_style=False,
|
||||
dtype=dtype,
|
||||
mrope_section=[8, 12, 12],
|
||||
).to(device)
|
||||
|
||||
torch.manual_seed(0)
|
||||
multimodal_positions = torch.randint(
|
||||
0,
|
||||
64,
|
||||
(3, 2 * num_tokens),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)[:, ::2]
|
||||
if position_kind == "multimodal":
|
||||
positions = multimodal_positions
|
||||
reference_positions = positions
|
||||
else:
|
||||
positions = torch.randint(
|
||||
0,
|
||||
64,
|
||||
(2 * num_tokens,),
|
||||
dtype=torch.int64,
|
||||
device=device,
|
||||
)[::2]
|
||||
reference_positions = positions.expand(3, -1)
|
||||
query_buffer = torch.randn(
|
||||
num_tokens,
|
||||
num_query_heads,
|
||||
192,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
key_buffer = torch.randn(
|
||||
num_tokens,
|
||||
576,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
query = query_buffer[..., 128:]
|
||||
key = key_buffer[:, 512:].unsqueeze(1)
|
||||
assert not query.is_contiguous()
|
||||
assert not key.is_contiguous()
|
||||
|
||||
query_prefix = query_buffer[..., :128].clone()
|
||||
key_prefix = key_buffer[:, :512].clone()
|
||||
expected_query, expected_key = rotary_emb.forward_native(
|
||||
reference_positions,
|
||||
query.clone(),
|
||||
key.clone(),
|
||||
)
|
||||
actual_query, actual_key = rotary_emb.forward_cuda(positions, query, key)
|
||||
assert expected_key is not None
|
||||
assert actual_key is not None
|
||||
|
||||
torch.testing.assert_close(actual_query, expected_query, atol=1e-2, rtol=1.6e-2)
|
||||
torch.testing.assert_close(actual_key, expected_key, atol=1e-2, rtol=1.6e-2)
|
||||
assert torch.equal(query_buffer[..., :128], query_prefix)
|
||||
assert torch.equal(key_buffer[:, :512], key_prefix)
|
||||
if position_kind == "multimodal":
|
||||
assert actual_query.data_ptr() == query.data_ptr()
|
||||
assert actual_key.data_ptr() == key.data_ptr()
|
||||
@@ -0,0 +1,126 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Bailing prompt expansion tests without downloading a checkpoint."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
from tokenizers import Tokenizer, decoders, models, pre_tokenizers
|
||||
from transformers import PreTrainedTokenizerFast
|
||||
|
||||
from vllm.model_executor.models.bailing_moe_v3_vl import (
|
||||
IMAGE_PLACEHOLDER,
|
||||
IMAGE_TOKEN,
|
||||
BailingMoeV3VLMultiModalProcessor,
|
||||
)
|
||||
from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
|
||||
from vllm.multimodal.parse import ImageProcessorItems, MultiModalDataItems
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def processor():
|
||||
vocab = {
|
||||
token: index
|
||||
for index, token in enumerate(sorted(pre_tokenizers.ByteLevel.alphabet()))
|
||||
}
|
||||
# A newline added by the processor must merge with an existing newline.
|
||||
vocab["ĊĊ"] = len(vocab)
|
||||
backend = Tokenizer(models.BPE(vocab=vocab, merges=[("Ċ", "Ċ")]))
|
||||
backend.pre_tokenizer = pre_tokenizers.ByteLevel(
|
||||
add_prefix_space=False, use_regex=False
|
||||
)
|
||||
backend.decoder = decoders.ByteLevel()
|
||||
tokenizer = PreTrainedTokenizerFast(
|
||||
tokenizer_object=backend,
|
||||
additional_special_tokens=[
|
||||
"<|vision_start|>",
|
||||
IMAGE_TOKEN,
|
||||
"<|vision_end|>",
|
||||
],
|
||||
clean_up_tokenization_spaces=False,
|
||||
)
|
||||
processor = object.__new__(BailingMoeV3VLMultiModalProcessor)
|
||||
processor.info = SimpleNamespace(
|
||||
get_tokenizer=lambda: tokenizer,
|
||||
get_hf_config=lambda: SimpleNamespace(
|
||||
vision_config=SimpleNamespace(spatial_merge_size=2)
|
||||
),
|
||||
)
|
||||
return processor
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("prompt", "expected_text", "token_counts"),
|
||||
[
|
||||
(
|
||||
f"before{IMAGE_PLACEHOLDER}after",
|
||||
f"before<|vision_start|>{IMAGE_TOKEN * 2}<|vision_end|>\nafter",
|
||||
[2],
|
||||
),
|
||||
(
|
||||
f"before{IMAGE_TOKEN}after",
|
||||
f"before<|vision_start|>{IMAGE_TOKEN * 2}<|vision_end|>\nafter",
|
||||
[2],
|
||||
),
|
||||
(
|
||||
f"before{IMAGE_PLACEHOLDER}\nafter",
|
||||
f"before<|vision_start|>{IMAGE_TOKEN * 2}<|vision_end|>\n\nafter",
|
||||
[2],
|
||||
),
|
||||
(
|
||||
f"before{IMAGE_PLACEHOLDER}{IMAGE_PLACEHOLDER}after",
|
||||
f"before<|vision_start|>{IMAGE_TOKEN * 2}<|vision_end|>\n"
|
||||
f"<|vision_start|>{IMAGE_TOKEN * 3}<|vision_end|>\nafter",
|
||||
[2, 3],
|
||||
),
|
||||
],
|
||||
ids=["wrapped", "bare", "existing_newline", "adjacent_images"],
|
||||
)
|
||||
def test_image_expansion_matches_wrappers_and_newline_tokenization(
|
||||
processor, prompt, expected_text, token_counts
|
||||
):
|
||||
tokenizer = processor.info.get_tokenizer()
|
||||
image_id = tokenizer.get_vocab()[IMAGE_TOKEN]
|
||||
grids = torch.tensor([[1, 2, 2 * count] for count in token_counts])
|
||||
mm_kwargs = MultiModalKwargsItems.from_hf_inputs(
|
||||
{"image_grid_thw": grids},
|
||||
{"image_grid_thw": MultiModalFieldConfig.batched("image")},
|
||||
)
|
||||
mm_items = MultiModalDataItems(
|
||||
{"image": ImageProcessorItems([Image.new("RGB", (1, 1))] * len(token_counts))}
|
||||
)
|
||||
updates = processor._get_mm_prompt_updates(mm_items, {}, mm_kwargs)
|
||||
prompt_ids = tokenizer.encode(prompt, add_special_tokens=False)
|
||||
expected_ids = tokenizer.encode(expected_text, add_special_tokens=False)
|
||||
|
||||
actual_ids, placeholders = processor._apply_prompt_updates(prompt_ids, updates)
|
||||
|
||||
assert actual_ids == expected_ids
|
||||
assert len(placeholders["image"]) == len(token_counts)
|
||||
found = processor._find_mm_placeholders(actual_ids, updates)
|
||||
search_offset = 0
|
||||
for index, count in enumerate(token_counts):
|
||||
placeholder = placeholders["image"][index]
|
||||
expected_offset = expected_ids.index(image_id, search_offset)
|
||||
assert placeholder.start_idx == expected_offset
|
||||
assert placeholder.tokens == [image_id] * count
|
||||
assert placeholder.to_range().get_num_embeds() == count
|
||||
assert placeholder.to_range() == found["image"][index].to_range()
|
||||
search_offset = expected_offset + count
|
||||
|
||||
|
||||
def test_text_only_prompt_preserves_original_token_ids(processor):
|
||||
tokenizer = processor.info.get_tokenizer()
|
||||
prompt_ids = tokenizer.encode("before\n", add_special_tokens=False)
|
||||
prompt_ids += tokenizer.encode("\nafter", add_special_tokens=False)
|
||||
# Deliberately use two newline tokens instead of their merged BPE token.
|
||||
assert prompt_ids != tokenizer.encode("before\n\nafter", add_special_tokens=False)
|
||||
|
||||
actual_ids, placeholders = processor._apply_prompt_updates(prompt_ids, {})
|
||||
|
||||
assert actual_ids == prompt_ids
|
||||
assert placeholders == {}
|
||||
@@ -184,6 +184,36 @@ def test_cosmos3_edge_checkpoint_weights_mapper():
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.cpu_test
|
||||
def test_bailing_vl_mapper_handles_module_and_parameter_names():
|
||||
from vllm.model_executor.models.bailing_moe_v3_vl import (
|
||||
BailingMoeV3VLForConditionalGeneration,
|
||||
)
|
||||
|
||||
mapper = BailingMoeV3VLForConditionalGeneration.hf_to_vllm_mapper
|
||||
names = {
|
||||
"model.visual": "visual",
|
||||
"model.visual.blocks.0.attn.qkv.weight": "visual.blocks.0.attn.qkv.weight",
|
||||
"model.linear_proj": "linear_proj",
|
||||
"linear_proj.0": "linear_proj.linear_fc1",
|
||||
"model.linear_proj.0.weight": "linear_proj.linear_fc1.weight",
|
||||
"linear_proj.2": "linear_proj.linear_fc2",
|
||||
"linear_proj.2.weight_scale": "linear_proj.linear_fc2.weight_scale",
|
||||
"model.linear_proj.2.bias": "linear_proj.linear_fc2.bias",
|
||||
"linear_proj.linear_fc1.weight": "linear_proj.linear_fc1.weight",
|
||||
"model.linear_proj.linear_fc2.weight": "linear_proj.linear_fc2.weight",
|
||||
"model.layers.0.attention.q_proj.weight": (
|
||||
"language_model.model.layers.0.self_attn.q_proj.weight"
|
||||
),
|
||||
"model.layers.0.mlp.experts.3.up_proj.weight": (
|
||||
"language_model.model.layers.0.mlp.experts.3.up_proj.weight"
|
||||
),
|
||||
"lm_head": "language_model.lm_head",
|
||||
"lm_head.weight": "language_model.lm_head.weight",
|
||||
}
|
||||
assert mapper.apply_list(list(names)) == list(names.values())
|
||||
|
||||
|
||||
def create_repo_dummy_weights(repo: str) -> Iterable[tuple[str, torch.Tensor]]:
|
||||
"""Create weights from safetensors checkpoint metadata"""
|
||||
metadata = try_get_safetensors_metadata(repo)
|
||||
|
||||
@@ -796,6 +796,18 @@ _MULTIMODAL_EXAMPLE_MODELS = {
|
||||
},
|
||||
),
|
||||
"BagelForConditionalGeneration": _HfExamplesInfo("ByteDance-Seed/BAGEL-7B-MoT"),
|
||||
"BailingMoeV3VLForConditionalGeneration": _HfExamplesInfo(
|
||||
"inclusionAI/Ling-3.0-flash-VL",
|
||||
trust_remote_code=True,
|
||||
is_available_online=True,
|
||||
use_original_num_layers=True,
|
||||
hf_overrides={
|
||||
"text_config": {
|
||||
"num_hidden_layers": 6,
|
||||
"layer_types": ["linear_attention"] * 5 + ["full_attention"],
|
||||
}
|
||||
},
|
||||
),
|
||||
"BeeForConditionalGeneration": _HfExamplesInfo(
|
||||
"Open-Bee/Bee-8B-RL",
|
||||
trust_remote_code=True,
|
||||
|
||||
@@ -3,6 +3,8 @@
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.parser.engine.replay_harness import MockTokenizer, _test_request
|
||||
from vllm.reasoning import ReasoningParserManager
|
||||
from vllm.tool_parsers import ToolParserManager
|
||||
@@ -57,19 +59,18 @@ def test_ling3_disable_thinking_keeps_reasoning_as_content():
|
||||
assert content == "<think>reason</think>answer"
|
||||
|
||||
|
||||
def test_ling3_enable_thinking_keeps_open_reasoning_as_content():
|
||||
@pytest.mark.parametrize("model_output", ["only reasoning", "<think>only reasoning"])
|
||||
def test_ling3_truncated_reasoning_stays_reasoning(model_output):
|
||||
parser_cls = ReasoningParserManager.get_reasoning_parser("ling3")
|
||||
parser = parser_cls(
|
||||
_tokenizer(),
|
||||
chat_template_kwargs={"enable_thinking": True},
|
||||
)
|
||||
|
||||
reasoning, content = parser.extract_reasoning(
|
||||
"<think>only reasoning", _test_request()
|
||||
)
|
||||
reasoning, content = parser.extract_reasoning(model_output, _test_request())
|
||||
|
||||
assert reasoning is None
|
||||
assert content == "only reasoning"
|
||||
assert reasoning == "only reasoning"
|
||||
assert content is None
|
||||
|
||||
|
||||
def test_ling3_tool_call_without_newline():
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import copy
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from transformers import AutoConfig
|
||||
|
||||
from vllm.transformers_utils.config import get_config, uses_mrope
|
||||
from vllm.transformers_utils.configs.bailing_moe_v3_vl import (
|
||||
BailingMoeV3TextConfig,
|
||||
BailingMoeV3VisionConfig,
|
||||
BailingMoeV3VLConfig,
|
||||
)
|
||||
|
||||
|
||||
def _bailing_v3_vl_config() -> dict[str, Any]:
|
||||
return {
|
||||
"architectures": ["BailingMoeV3VLForConditionalGeneration"],
|
||||
"auto_map": {
|
||||
"AutoConfig": ("configuration_bailing_moe_v3_vl.BailingMoeV3VLConfig")
|
||||
},
|
||||
"model_type": "bailing_moe_v3_vl",
|
||||
"mrope_section": [8, 12, 12],
|
||||
"text_config": {
|
||||
"num_hidden_layers": 42,
|
||||
"layer_group_size": 6,
|
||||
"rope_theta": 6_000_000,
|
||||
},
|
||||
"vision_config": {"disable_merger_proj": True},
|
||||
}
|
||||
|
||||
|
||||
def _write_config(path: Path, config: dict[str, Any]) -> None:
|
||||
path.mkdir()
|
||||
(path / "config.json").write_text(json.dumps(config), encoding="utf-8")
|
||||
|
||||
|
||||
def test_bailing_v3_vl_config_loads_without_remote_code(tmp_path: Path):
|
||||
model_path = tmp_path / "model"
|
||||
_write_config(model_path, _bailing_v3_vl_config())
|
||||
|
||||
config = get_config(model_path, trust_remote_code=False)
|
||||
|
||||
assert isinstance(config, BailingMoeV3VLConfig)
|
||||
assert isinstance(config.text_config, BailingMoeV3TextConfig)
|
||||
assert isinstance(config.vision_config, BailingMoeV3VisionConfig)
|
||||
assert isinstance(
|
||||
AutoConfig.from_pretrained(model_path, trust_remote_code=False),
|
||||
BailingMoeV3VLConfig,
|
||||
)
|
||||
assert config.mrope_section == [8, 12, 12]
|
||||
assert uses_mrope(config)
|
||||
|
||||
text_config = config.text_config
|
||||
assert text_config.model_type == "bailing_hybrid"
|
||||
assert text_config.architectures == ["BailingMoeV3ForCausalLM"]
|
||||
assert text_config.rope_parameters["mrope_section"] == [8, 12, 12]
|
||||
assert text_config.rope_parameters["rope_theta"] == 6_000_000
|
||||
assert text_config.layer_types.count("linear_attention") == 35
|
||||
assert [
|
||||
layer_idx
|
||||
for layer_idx, layer_type in enumerate(text_config.layer_types)
|
||||
if layer_type == "full_attention"
|
||||
] == [5, 11, 17, 23, 29, 35, 41]
|
||||
assert config.vision_config.disable_merger_proj is True
|
||||
|
||||
normalized = config.to_dict()
|
||||
restored = BailingMoeV3VLConfig.from_dict(normalized)
|
||||
assert isinstance(restored.text_config, BailingMoeV3TextConfig)
|
||||
assert isinstance(restored.vision_config, BailingMoeV3VisionConfig)
|
||||
assert restored.to_dict() == normalized
|
||||
|
||||
|
||||
def test_bailing_v3_vl_config_preserves_explicit_nested_values():
|
||||
config_dict = _bailing_v3_vl_config()
|
||||
config_dict["architectures"] = ["CustomBailingVLForConditionalGeneration"]
|
||||
text_config = config_dict["text_config"]
|
||||
text_config["architectures"] = ["CustomBailingForCausalLM"]
|
||||
text_config["layer_types"] = (["linear_attention"] * 5 + ["full_attention"]) * 7
|
||||
text_config["rope_theta"] = 1_000_000
|
||||
text_config["rope_parameters"] = {
|
||||
"rope_type": "default",
|
||||
"rope_theta": 1_000_000,
|
||||
"mrope_section": [8, 12, 12],
|
||||
}
|
||||
original = copy.deepcopy(config_dict)
|
||||
|
||||
config = BailingMoeV3VLConfig(**config_dict)
|
||||
|
||||
assert config_dict == original
|
||||
assert config.architectures == ["CustomBailingVLForConditionalGeneration"]
|
||||
assert config.text_config.architectures == ["CustomBailingForCausalLM"]
|
||||
assert config.text_config.layer_types == text_config["layer_types"]
|
||||
assert config.text_config.rope_parameters["rope_theta"] == 1_000_000
|
||||
assert config.text_config.rope_parameters["mrope_section"] == [8, 12, 12]
|
||||
|
||||
|
||||
def test_bailing_v3_vl_config_normalizes_config_instance():
|
||||
text_config = BailingMoeV3TextConfig()
|
||||
|
||||
config = BailingMoeV3VLConfig(
|
||||
text_config=text_config,
|
||||
mrope_section=[8, 12, 12],
|
||||
)
|
||||
|
||||
assert config.text_config is text_config
|
||||
assert text_config.rope_parameters["mrope_section"] == [8, 12, 12]
|
||||
|
||||
|
||||
def test_bailing_v3_text_config_trailing_and_legacy_layers():
|
||||
expected = ["linear_attention", "linear_attention", "full_attention"] * 2 + [
|
||||
"full_attention"
|
||||
] * 2
|
||||
text_config = BailingMoeV3TextConfig(
|
||||
num_hidden_layers=8,
|
||||
layer_group_size=3,
|
||||
)
|
||||
assert text_config.layer_types == expected
|
||||
|
||||
text_config = BailingMoeV3TextConfig(
|
||||
num_hidden_layers=8,
|
||||
layer_group_size=3,
|
||||
layers_block_type=["mamba", "mamba", "attention"] * 2 + ["attention"] * 2,
|
||||
)
|
||||
assert text_config.layer_types == expected
|
||||
|
||||
|
||||
def test_bailing_v3_text_config_accepts_legacy_rope_type():
|
||||
text_config = BailingMoeV3TextConfig(rope_scaling={"type": "linear", "factor": 2.0})
|
||||
|
||||
assert text_config.rope_parameters["type"] == "linear"
|
||||
assert text_config.rope_parameters["rope_type"] == "linear"
|
||||
|
||||
|
||||
def test_bailing_v3_vl_config_rejects_invalid_values():
|
||||
with pytest.raises(ValueError, match="layer_group_size must be positive"):
|
||||
BailingMoeV3TextConfig(layer_group_size=0)
|
||||
|
||||
with pytest.raises(ValueError, match="must match the layer_group_size schedule"):
|
||||
BailingMoeV3TextConfig(
|
||||
num_hidden_layers=8,
|
||||
layer_group_size=3,
|
||||
layer_types=["full_attention"] * 8,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="conflicting M-RoPE sections"):
|
||||
BailingMoeV3VLConfig(
|
||||
text_config={
|
||||
"rope_parameters": {
|
||||
"rope_type": "default",
|
||||
"mrope_section": [4, 14, 14],
|
||||
}
|
||||
},
|
||||
mrope_section=[8, 12, 12],
|
||||
)
|
||||
@@ -847,7 +847,18 @@ class SpeculativeConfig:
|
||||
)
|
||||
|
||||
architectures = getattr(hf_config, "architectures", []) or []
|
||||
if initial_architecture == "BailingMoeV3ForCausalLM":
|
||||
if initial_architecture == "BailingMoeV3VLForConditionalGeneration":
|
||||
quantization_config = getattr(hf_config, "quantization_config", None)
|
||||
hf_config = copy.deepcopy(hf_config.text_config)
|
||||
if (
|
||||
quantization_config is not None
|
||||
and getattr(hf_config, "quantization_config", None) is None
|
||||
):
|
||||
hf_config.quantization_config = copy.deepcopy(quantization_config)
|
||||
if initial_architecture in (
|
||||
"BailingMoeV3ForCausalLM",
|
||||
"BailingMoeV3VLForConditionalGeneration",
|
||||
):
|
||||
hf_config.model_type = "bailing_hybrid_v3_mtp"
|
||||
elif (
|
||||
hf_config.model_type == "bailing_hybrid"
|
||||
|
||||
@@ -6,6 +6,7 @@ from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from .bailing_mrope import BailingMRotaryEmbedding
|
||||
from .base import RotaryEmbedding
|
||||
from .deepseek_scaling_rope import (
|
||||
DeepseekScalingRotaryEmbedding,
|
||||
@@ -97,6 +98,16 @@ def get_rope(
|
||||
dtype,
|
||||
**extra_kwargs,
|
||||
)
|
||||
elif scaling_type == "bailing_mrope":
|
||||
rotary_emb = BailingMRotaryEmbedding(
|
||||
head_size,
|
||||
rotary_dim,
|
||||
max_position,
|
||||
base,
|
||||
is_neox_style,
|
||||
dtype,
|
||||
mrope_section=rope_parameters["mrope_section"],
|
||||
)
|
||||
elif scaling_type == "default":
|
||||
if "mrope_section" in rope_parameters:
|
||||
rotary_emb = MRotaryEmbedding(
|
||||
|
||||
@@ -0,0 +1,351 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.triton_utils import tl, triton
|
||||
|
||||
from .mrope import MRotaryEmbedding
|
||||
|
||||
|
||||
def _validate_bailing_mrope_config(
|
||||
head_size: int,
|
||||
rotary_dim: int,
|
||||
mrope_section: list[int],
|
||||
) -> int:
|
||||
if rotary_dim <= 0 or rotary_dim % 2 != 0 or rotary_dim > head_size:
|
||||
raise ValueError(
|
||||
"Bailing M-RoPE requires an even rotary_dim in (0, head_size], "
|
||||
f"got rotary_dim={rotary_dim} and head_size={head_size}"
|
||||
)
|
||||
if len(mrope_section) != 3:
|
||||
raise ValueError(
|
||||
"Bailing M-RoPE expects [T, H, W] sections, "
|
||||
f"got mrope_section={mrope_section}"
|
||||
)
|
||||
|
||||
temporal_size, height_size, width_size = mrope_section
|
||||
if any(section_size < 0 for section_size in mrope_section):
|
||||
raise ValueError(
|
||||
"Bailing M-RoPE sections must be non-negative, "
|
||||
f"got mrope_section={mrope_section}"
|
||||
)
|
||||
if height_size != width_size:
|
||||
raise ValueError(
|
||||
"Bailing M-RoPE requires equal H and W sections, "
|
||||
f"got height_size={height_size} and width_size={width_size}"
|
||||
)
|
||||
spatial_size = height_size + width_size
|
||||
if spatial_size + temporal_size != rotary_dim // 2:
|
||||
raise ValueError(
|
||||
"Bailing M-RoPE sections must sum to rotary_dim // 2, got "
|
||||
f"mrope_section={mrope_section} and rotary_dim={rotary_dim}"
|
||||
)
|
||||
return spatial_size
|
||||
|
||||
|
||||
def _validate_bailing_positions(
|
||||
positions: torch.Tensor,
|
||||
*,
|
||||
require_2d: bool = False,
|
||||
) -> None:
|
||||
if positions.dtype not in (torch.int32, torch.int64):
|
||||
raise ValueError(
|
||||
"Bailing M-RoPE positions must have int32 or int64 dtype, got "
|
||||
f"dtype={positions.dtype}"
|
||||
)
|
||||
if positions.ndim == 1 and not require_2d:
|
||||
return
|
||||
if positions.ndim == 2:
|
||||
if positions.shape[0] == 3:
|
||||
return
|
||||
raise ValueError(
|
||||
"Bailing M-RoPE expects T/H/W positions with shape [3, N], "
|
||||
f"got shape={tuple(positions.shape)}"
|
||||
)
|
||||
if require_2d:
|
||||
raise ValueError(
|
||||
"Bailing M-RoPE expects T/H/W positions with shape [3, N], "
|
||||
f"got shape={tuple(positions.shape)}"
|
||||
)
|
||||
raise ValueError(
|
||||
f"Bailing M-RoPE positions must be 1D or 2D, got shape={tuple(positions.shape)}"
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _triton_bailing_mrope_forward(
|
||||
q_ptr,
|
||||
k_ptr,
|
||||
positions_ptr,
|
||||
cos_sin_cache_ptr,
|
||||
q_stride_token,
|
||||
q_stride_head,
|
||||
q_stride_dim,
|
||||
k_stride_token,
|
||||
k_stride_head,
|
||||
k_stride_dim,
|
||||
positions_stride_axis,
|
||||
positions_stride_token,
|
||||
cache_stride_position,
|
||||
cache_stride_dim,
|
||||
n_qh: tl.constexpr,
|
||||
n_kh: tl.constexpr,
|
||||
rotary_dim: tl.constexpr,
|
||||
pad_n_qh: tl.constexpr,
|
||||
pad_n_kh: tl.constexpr,
|
||||
pad_half_rotary_dim: tl.constexpr,
|
||||
spatial_size: tl.constexpr,
|
||||
):
|
||||
"""Apply GPT-J-style H/W-interleaved, T-tail M-RoPE in place."""
|
||||
token_idx = tl.program_id(0)
|
||||
position_offset = token_idx * positions_stride_token
|
||||
temporal_position = tl.load(positions_ptr + position_offset).to(tl.int64)
|
||||
height_position = tl.load(
|
||||
positions_ptr + positions_stride_axis + position_offset
|
||||
).to(tl.int64)
|
||||
width_position = tl.load(
|
||||
positions_ptr + 2 * positions_stride_axis + position_offset
|
||||
).to(tl.int64)
|
||||
|
||||
frequency_offsets = tl.arange(0, pad_half_rotary_dim)
|
||||
frequency_mask = frequency_offsets < rotary_dim // 2
|
||||
spatial_position = tl.where(
|
||||
frequency_offsets % 2 == 0,
|
||||
height_position,
|
||||
width_position,
|
||||
)
|
||||
selected_position = tl.where(
|
||||
frequency_offsets < spatial_size,
|
||||
spatial_position,
|
||||
temporal_position,
|
||||
)
|
||||
cache_offsets = (
|
||||
selected_position * cache_stride_position + frequency_offsets * cache_stride_dim
|
||||
)
|
||||
cos_row = tl.load(
|
||||
cos_sin_cache_ptr + cache_offsets,
|
||||
mask=frequency_mask,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
sin_row = tl.load(
|
||||
cos_sin_cache_ptr + cache_offsets + (rotary_dim // 2) * cache_stride_dim,
|
||||
mask=frequency_mask,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
|
||||
q_head_offsets = tl.arange(0, pad_n_qh)[:, None]
|
||||
k_head_offsets = tl.arange(0, pad_n_kh)[:, None]
|
||||
pair_offsets = frequency_offsets[None, :]
|
||||
q_mask = (q_head_offsets < n_qh) & frequency_mask[None, :]
|
||||
k_mask = (k_head_offsets < n_kh) & frequency_mask[None, :]
|
||||
|
||||
q_base = q_ptr + token_idx * q_stride_token
|
||||
q_even_offsets = q_head_offsets * q_stride_head + 2 * pair_offsets * q_stride_dim
|
||||
q_odd_offsets = q_even_offsets + q_stride_dim
|
||||
q_even = tl.load(q_base + q_even_offsets, mask=q_mask, other=0.0).to(tl.float32)
|
||||
q_odd = tl.load(q_base + q_odd_offsets, mask=q_mask, other=0.0).to(tl.float32)
|
||||
q_even_cos = (q_even * cos_row[None, :]).to(q_ptr.type.element_ty).to(tl.float32)
|
||||
q_odd_sin = (q_odd * sin_row[None, :]).to(q_ptr.type.element_ty).to(tl.float32)
|
||||
q_odd_cos = (q_odd * cos_row[None, :]).to(q_ptr.type.element_ty).to(tl.float32)
|
||||
q_even_sin = (q_even * sin_row[None, :]).to(q_ptr.type.element_ty).to(tl.float32)
|
||||
tl.store(
|
||||
q_base + q_even_offsets,
|
||||
q_even_cos - q_odd_sin,
|
||||
mask=q_mask,
|
||||
)
|
||||
tl.store(
|
||||
q_base + q_odd_offsets,
|
||||
q_odd_cos + q_even_sin,
|
||||
mask=q_mask,
|
||||
)
|
||||
|
||||
k_base = k_ptr + token_idx * k_stride_token
|
||||
k_even_offsets = k_head_offsets * k_stride_head + 2 * pair_offsets * k_stride_dim
|
||||
k_odd_offsets = k_even_offsets + k_stride_dim
|
||||
k_even = tl.load(k_base + k_even_offsets, mask=k_mask, other=0.0).to(tl.float32)
|
||||
k_odd = tl.load(k_base + k_odd_offsets, mask=k_mask, other=0.0).to(tl.float32)
|
||||
k_even_cos = (k_even * cos_row[None, :]).to(k_ptr.type.element_ty).to(tl.float32)
|
||||
k_odd_sin = (k_odd * sin_row[None, :]).to(k_ptr.type.element_ty).to(tl.float32)
|
||||
k_odd_cos = (k_odd * cos_row[None, :]).to(k_ptr.type.element_ty).to(tl.float32)
|
||||
k_even_sin = (k_even * sin_row[None, :]).to(k_ptr.type.element_ty).to(tl.float32)
|
||||
tl.store(
|
||||
k_base + k_even_offsets,
|
||||
k_even_cos - k_odd_sin,
|
||||
mask=k_mask,
|
||||
)
|
||||
tl.store(
|
||||
k_base + k_odd_offsets,
|
||||
k_odd_cos + k_even_sin,
|
||||
mask=k_mask,
|
||||
)
|
||||
|
||||
|
||||
def triton_bailing_mrope(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
spatial_size: int,
|
||||
head_size: int,
|
||||
rotary_dim: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Fused Bailing M-RoPE for packed or strided Q/K."""
|
||||
_validate_bailing_positions(positions, require_2d=True)
|
||||
num_tokens = positions.shape[-1]
|
||||
if num_tokens == 0:
|
||||
return query, key
|
||||
|
||||
query_3d = query.view(num_tokens, -1, head_size)
|
||||
key_3d = key.view(num_tokens, -1, head_size)
|
||||
n_q_heads = query_3d.shape[1]
|
||||
n_k_heads = key_3d.shape[1]
|
||||
pad_n_q_heads = triton.next_power_of_2(n_q_heads)
|
||||
pad_n_k_heads = triton.next_power_of_2(n_k_heads)
|
||||
pad_half_rotary_dim = triton.next_power_of_2(rotary_dim // 2)
|
||||
_triton_bailing_mrope_forward[(num_tokens,)](
|
||||
query_3d,
|
||||
key_3d,
|
||||
positions,
|
||||
cos_sin_cache,
|
||||
query_3d.stride(0),
|
||||
query_3d.stride(1),
|
||||
query_3d.stride(2),
|
||||
key_3d.stride(0),
|
||||
key_3d.stride(1),
|
||||
key_3d.stride(2),
|
||||
positions.stride(0),
|
||||
positions.stride(1),
|
||||
cos_sin_cache.stride(0),
|
||||
cos_sin_cache.stride(1),
|
||||
n_q_heads,
|
||||
n_k_heads,
|
||||
rotary_dim,
|
||||
pad_n_q_heads,
|
||||
pad_n_k_heads,
|
||||
pad_half_rotary_dim,
|
||||
spatial_size,
|
||||
num_warps=4,
|
||||
)
|
||||
return query, key
|
||||
|
||||
|
||||
class BailingMRotaryEmbedding(MRotaryEmbedding):
|
||||
"""Bailing M-RoPE with alternating H/W frequencies and a temporal tail."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
head_size: int,
|
||||
rotary_dim: int,
|
||||
max_position_embeddings: int,
|
||||
base: float,
|
||||
is_neox_style: bool,
|
||||
dtype: torch.dtype,
|
||||
mrope_section: list[int],
|
||||
) -> None:
|
||||
self.spatial_size = _validate_bailing_mrope_config(
|
||||
head_size, rotary_dim, mrope_section
|
||||
)
|
||||
super().__init__(
|
||||
head_size,
|
||||
rotary_dim,
|
||||
max_position_embeddings,
|
||||
base,
|
||||
is_neox_style,
|
||||
dtype,
|
||||
mrope_section=mrope_section,
|
||||
)
|
||||
|
||||
def select_cos_sin(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
_validate_bailing_positions(positions)
|
||||
if cos_sin_cache is None:
|
||||
cos_sin_cache = self.cos_sin_cache
|
||||
cos_sin = cos_sin_cache[positions.to(torch.long)]
|
||||
cos, sin = cos_sin.chunk(2, dim=-1)
|
||||
if positions.ndim == 1:
|
||||
return cos, sin
|
||||
|
||||
spatial_mask = torch.arange(self.spatial_size, device=cos.device) % 2 == 0
|
||||
spatial_mask = spatial_mask.unsqueeze(0)
|
||||
spatial_cos = torch.where(
|
||||
spatial_mask,
|
||||
cos[1, :, : self.spatial_size],
|
||||
cos[2, :, : self.spatial_size],
|
||||
)
|
||||
spatial_sin = torch.where(
|
||||
spatial_mask,
|
||||
sin[1, :, : self.spatial_size],
|
||||
sin[2, :, : self.spatial_size],
|
||||
)
|
||||
temporal_slice = slice(self.spatial_size, self.rotary_dim // 2)
|
||||
return (
|
||||
torch.cat((spatial_cos, cos[0, :, temporal_slice]), dim=-1),
|
||||
torch.cat((spatial_sin, sin[0, :, temporal_slice]), dim=-1),
|
||||
)
|
||||
|
||||
def forward_native(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor | None = None,
|
||||
offsets: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
if key is None:
|
||||
raise ValueError("Bailing M-RoPE requires key states")
|
||||
|
||||
cos_sin_cache = self._match_cos_sin_cache_dtype(query)
|
||||
cos, sin = self.select_cos_sin(positions, cos_sin_cache)
|
||||
num_tokens = positions.shape[-1]
|
||||
query_shape = query.shape
|
||||
query = query.view(num_tokens, -1, self.head_size)
|
||||
query_rot = self.apply_rotary_emb.forward_native(
|
||||
query[..., : self.rotary_dim],
|
||||
cos,
|
||||
sin,
|
||||
)
|
||||
query = torch.cat((query_rot, query[..., self.rotary_dim :]), dim=-1)
|
||||
|
||||
key_shape = key.shape
|
||||
key = key.view(num_tokens, -1, self.head_size)
|
||||
key_rot = self.apply_rotary_emb.forward_native(
|
||||
key[..., : self.rotary_dim],
|
||||
cos,
|
||||
sin,
|
||||
)
|
||||
key = torch.cat((key_rot, key[..., self.rotary_dim :]), dim=-1)
|
||||
return query.reshape(query_shape), key.reshape(key_shape)
|
||||
|
||||
def forward_cuda(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor | None = None,
|
||||
offsets: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
if positions.ndim == 2 and not self.is_neox_style:
|
||||
if key is None:
|
||||
raise ValueError("Bailing M-RoPE requires key states")
|
||||
cos_sin_cache = self._match_cos_sin_cache_dtype(query)
|
||||
return triton_bailing_mrope(
|
||||
query,
|
||||
key,
|
||||
positions,
|
||||
cos_sin_cache,
|
||||
self.spatial_size,
|
||||
self.head_size,
|
||||
self.rotary_dim,
|
||||
)
|
||||
return self.forward_native(positions, query, key, offsets)
|
||||
|
||||
def forward_cpu(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor | None = None,
|
||||
offsets: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
return self.forward_native(positions, query, key, offsets)
|
||||
@@ -171,6 +171,27 @@ def _build_rope_parameters(config: PretrainedConfig) -> dict | None:
|
||||
return rope_parameters or None
|
||||
|
||||
|
||||
def _build_mla_rotary_embedding(
|
||||
config: PretrainedConfig,
|
||||
head_size: int,
|
||||
) -> nn.Module:
|
||||
rope_parameters = _build_rope_parameters(config)
|
||||
if rope_parameters is not None and "mrope_section" in rope_parameters:
|
||||
rope_type = rope_parameters.get("rope_type", "default")
|
||||
if rope_type != "default":
|
||||
raise ValueError(
|
||||
f"Bailing M-RoPE only supports rope_type='default', got {rope_type!r}"
|
||||
)
|
||||
rope_parameters["rope_type"] = "bailing_mrope"
|
||||
|
||||
return get_rope(
|
||||
head_size=head_size,
|
||||
max_position=getattr(config, "max_position_embeddings", 8192),
|
||||
is_neox_style=False,
|
||||
rope_parameters=rope_parameters,
|
||||
)
|
||||
|
||||
|
||||
def _get_layer_swiglu_limit(limit_list: list | None, layer_idx: int) -> float | None:
|
||||
if limit_list is None or layer_idx >= len(limit_list):
|
||||
return None
|
||||
@@ -524,11 +545,9 @@ class BailingMoeV3MLAAttention(nn.Module):
|
||||
prefix=f"{prefix}.dense",
|
||||
)
|
||||
|
||||
self.rotary_emb = get_rope(
|
||||
self.rotary_emb = _build_mla_rotary_embedding(
|
||||
config,
|
||||
head_size=self.qk_rope_head_dim,
|
||||
max_position=getattr(config, "max_position_embeddings", 8192),
|
||||
is_neox_style=False,
|
||||
rope_parameters=_build_rope_parameters(config),
|
||||
)
|
||||
mla_modules = MLAModules(
|
||||
kv_a_layernorm=self.kv_a_layernorm,
|
||||
|
||||
@@ -33,9 +33,14 @@ from vllm.model_executor.models.bailing_moe_v3 import (
|
||||
_maybe_pad_block_fp8_shared_expert_checkpoint_tensor,
|
||||
_maybe_remap_ling_mxfp4_weight_names,
|
||||
)
|
||||
from vllm.model_executor.models.interfaces import SupportsPP
|
||||
from vllm.model_executor.models.interfaces import (
|
||||
MultiModalEmbeddings,
|
||||
SupportsMultiModalEmbeddings,
|
||||
SupportsPP,
|
||||
)
|
||||
from vllm.model_executor.models.utils import (
|
||||
PPMissingLayer,
|
||||
_merge_multimodal_embeddings,
|
||||
is_pp_missing_parameter,
|
||||
maybe_prefix,
|
||||
)
|
||||
@@ -214,8 +219,16 @@ class BailingMoeV3MultiTokenPredictor(nn.Module):
|
||||
return self.logits_processor(head, mtp_layer.shared_head(hidden_states))
|
||||
|
||||
|
||||
@support_torch_compile
|
||||
class BailingMoeV3MTPModel(nn.Module, SupportsPP):
|
||||
@support_torch_compile(
|
||||
dynamic_arg_dims={
|
||||
"input_ids": 0,
|
||||
"positions": -1,
|
||||
"hidden_states": 0,
|
||||
"intermediate_tensors": 0,
|
||||
"inputs_embeds": 0,
|
||||
}
|
||||
)
|
||||
class BailingMoeV3MTPModel(nn.Module, SupportsPP, SupportsMultiModalEmbeddings):
|
||||
hf_to_vllm_mapper = BailingMoeV3ForCausalLM.hf_to_vllm_mapper
|
||||
|
||||
packed_modules_mapping = {
|
||||
@@ -242,8 +255,22 @@ class BailingMoeV3MTPModel(nn.Module, SupportsPP):
|
||||
def share_lm_head(self, lm_head: nn.Module) -> None:
|
||||
self.lm_head = lm_head
|
||||
|
||||
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
return self.model.embed_input_ids(input_ids)
|
||||
def embed_input_ids(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
multimodal_embeddings: MultiModalEmbeddings | None = None,
|
||||
*,
|
||||
is_multimodal: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
inputs_embeds = self.model.embed_input_ids(input_ids)
|
||||
if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
|
||||
return inputs_embeds
|
||||
assert is_multimodal is not None
|
||||
return _merge_multimodal_embeddings(
|
||||
inputs_embeds=inputs_embeds,
|
||||
multimodal_embeddings=multimodal_embeddings,
|
||||
is_multimodal=is_multimodal,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
|
||||
@@ -0,0 +1,527 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Inference-only Bailing MoE V3 vision-language model."""
|
||||
|
||||
import copy
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import regex as re
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from transformers import AutoProcessor, BatchFeature
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.config.multimodal import BaseDummyOptions
|
||||
from vllm.inputs import MultiModalDataDict
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization import QuantizationConfig
|
||||
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||||
from vllm.model_executor.models.module_mapping import MultiModelKeys
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
from vllm.multimodal.inputs import (
|
||||
MultiModalFeatureSpec,
|
||||
MultiModalFieldConfig,
|
||||
MultiModalKwargsItems,
|
||||
)
|
||||
from vllm.multimodal.parse import MultiModalDataItems
|
||||
from vllm.multimodal.processing import (
|
||||
BaseDummyInputsBuilder,
|
||||
BaseMultiModalProcessor,
|
||||
PromptReplacement,
|
||||
PromptUpdate,
|
||||
)
|
||||
from vllm.multimodal.processing.processor import (
|
||||
MultiModalPromptUpdates,
|
||||
PlaceholderFeaturesInfo,
|
||||
)
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
from .bailing_moe_v3 import BailingMoeV3ForCausalLM
|
||||
from .interfaces import (
|
||||
HasInnerState,
|
||||
IsHybrid,
|
||||
MultiModalEmbeddings,
|
||||
SupportsMRoPE,
|
||||
SupportsMultiModal,
|
||||
SupportsPP,
|
||||
_require_is_multimodal,
|
||||
)
|
||||
from .qwen2_5_vl import (
|
||||
Qwen2_5_VLImageEmbeddingInputs,
|
||||
Qwen2_5_VLImageInputs,
|
||||
Qwen2_5_VLImagePixelInputs,
|
||||
)
|
||||
from .qwen2_vl import (
|
||||
Qwen2VLProcessingInfo,
|
||||
_create_qwen2vl_field_factory,
|
||||
)
|
||||
from .qwen3_vl import Qwen3_VisionTransformer
|
||||
from .utils import (
|
||||
AutoWeightsLoader,
|
||||
WeightsMapper,
|
||||
_merge_multimodal_embeddings,
|
||||
maybe_prefix,
|
||||
)
|
||||
|
||||
IMAGE_TOKEN = "<|image_pad|>"
|
||||
IMAGE_PLACEHOLDER = f"<|vision_start|>{IMAGE_TOKEN}<|vision_end|>"
|
||||
|
||||
|
||||
class BailingMoeV3VLProcessingInfo(Qwen2VLProcessingInfo):
|
||||
"""Processing metadata for the remote Bailing multimodal processor."""
|
||||
|
||||
def get_hf_config(self):
|
||||
return self.ctx.get_hf_config()
|
||||
|
||||
def get_hf_processor(self, **kwargs: object) -> AutoProcessor:
|
||||
# The checkpoint supplies its processor through trust_remote_code.
|
||||
return self.ctx.get_hf_processor(**kwargs)
|
||||
|
||||
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
|
||||
return {"image": None}
|
||||
|
||||
def get_mm_max_tokens_per_item(
|
||||
self,
|
||||
seq_len: int,
|
||||
mm_counts: Mapping[str, int],
|
||||
) -> Mapping[str, int]:
|
||||
return {"image": self.get_max_image_tokens()}
|
||||
|
||||
|
||||
class BailingMoeV3VLDummyInputsBuilder(
|
||||
BaseDummyInputsBuilder[BailingMoeV3VLProcessingInfo]
|
||||
):
|
||||
def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
|
||||
return IMAGE_PLACEHOLDER * mm_counts.get("image", 0)
|
||||
|
||||
def get_dummy_mm_data(
|
||||
self,
|
||||
seq_len: int,
|
||||
mm_counts: Mapping[str, int],
|
||||
mm_options: Mapping[str, BaseDummyOptions],
|
||||
) -> MultiModalDataDict:
|
||||
target_width, target_height = self.info.get_image_size_with_most_features()
|
||||
return {
|
||||
"image": self._get_dummy_images(
|
||||
width=target_width,
|
||||
height=target_height,
|
||||
num_images=mm_counts.get("image", 0),
|
||||
overrides=mm_options.get("image"),
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
class BailingMoeV3VLMultiModalProcessor(
|
||||
BaseMultiModalProcessor[BailingMoeV3VLProcessingInfo]
|
||||
):
|
||||
def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
|
||||
return self.dummy_inputs.get_dummy_text(mm_counts)
|
||||
|
||||
def _apply_prompt_updates(
|
||||
self,
|
||||
token_ids: list[int],
|
||||
mm_prompt_updates: MultiModalPromptUpdates,
|
||||
) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
|
||||
if image_updates := mm_prompt_updates.get("image"):
|
||||
tokenizer = self.info.get_tokenizer()
|
||||
# HF adds a newline after each image before tokenization. Re-encode
|
||||
# the text so it can merge with existing whitespace at the boundary.
|
||||
text = re.sub(
|
||||
f"{re.escape(IMAGE_PLACEHOLDER)}|{re.escape(IMAGE_TOKEN)}",
|
||||
lambda _: IMAGE_PLACEHOLDER + "\n",
|
||||
tokenizer.decode(token_ids),
|
||||
count=len(image_updates),
|
||||
)
|
||||
token_ids = tokenizer.encode(text, add_special_tokens=False)
|
||||
|
||||
return super()._apply_prompt_updates(token_ids, mm_prompt_updates)
|
||||
|
||||
def _get_mm_fields_config(
|
||||
self,
|
||||
hf_inputs: BatchFeature,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
) -> Mapping[str, MultiModalFieldConfig]:
|
||||
merge_size = self.info.get_hf_config().vision_config.spatial_merge_size
|
||||
fields = _create_qwen2vl_field_factory(merge_size)(hf_inputs)
|
||||
return {
|
||||
key: fields[key]
|
||||
for key in ("pixel_values", "image_embeds", "image_grid_thw")
|
||||
}
|
||||
|
||||
def _get_prompt_updates(
|
||||
self,
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, Any],
|
||||
out_mm_kwargs: MultiModalKwargsItems,
|
||||
) -> Sequence[PromptUpdate]:
|
||||
tokenizer = self.info.get_tokenizer()
|
||||
image_token_id = tokenizer.get_vocab()[IMAGE_TOKEN]
|
||||
merge_size = self.info.get_hf_config().vision_config.spatial_merge_size
|
||||
merge_unit = merge_size**2
|
||||
|
||||
def get_image_replacement(item_idx: int) -> list[int]:
|
||||
out_item = out_mm_kwargs["image"][item_idx]
|
||||
image_grid_thw = out_item["image_grid_thw"].data
|
||||
assert isinstance(image_grid_thw, torch.Tensor)
|
||||
num_image_tokens = int(image_grid_thw.prod()) // merge_unit
|
||||
return [image_token_id] * num_image_tokens
|
||||
|
||||
return [
|
||||
PromptReplacement(
|
||||
modality="image",
|
||||
target=[image_token_id],
|
||||
replacement=get_image_replacement,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
class BailingMoeV3VisionPatchMerger(nn.Module):
|
||||
"""Normalize patches and merge each 2x2 group without an MLP."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
context_dim: int,
|
||||
spatial_merge_size: int,
|
||||
norm_eps: float = 1e-6,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.hidden_size = context_dim * spatial_merge_size**2
|
||||
self.norm = nn.LayerNorm(context_dim, eps=norm_eps)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
hidden_states = self.norm(hidden_states)
|
||||
return hidden_states.reshape(-1, self.hidden_size)
|
||||
|
||||
|
||||
class BailingMoeV3VisionTransformer(Qwen3_VisionTransformer):
|
||||
"""Qwen3 ViT with Bailing's parameter-free patch merger."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vision_config: Any,
|
||||
norm_eps: float = 1e-6,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
# Honor whole-tower exclusions before the parent creates linear layers.
|
||||
if (
|
||||
isinstance(quant_config, Fp8Config)
|
||||
and prefix in quant_config.ignored_layers
|
||||
):
|
||||
quant_config = None
|
||||
|
||||
vision_config = copy.deepcopy(vision_config)
|
||||
vision_config.deepstack_visual_indexes = []
|
||||
merged_hidden_size = (
|
||||
vision_config.hidden_size * vision_config.spatial_merge_size**2
|
||||
)
|
||||
vision_config.out_hidden_size = merged_hidden_size
|
||||
|
||||
super().__init__(
|
||||
vision_config,
|
||||
norm_eps=norm_eps,
|
||||
quant_config=quant_config,
|
||||
prefix=prefix,
|
||||
)
|
||||
self.merger = BailingMoeV3VisionPatchMerger(
|
||||
context_dim=vision_config.hidden_size,
|
||||
spatial_merge_size=vision_config.spatial_merge_size,
|
||||
norm_eps=norm_eps,
|
||||
)
|
||||
self.deepstack_visual_indexes = []
|
||||
self.deepstack_merger_list = nn.ModuleList()
|
||||
self.out_hidden_size = merged_hidden_size
|
||||
|
||||
|
||||
class BailingMoeV3VLProjector(nn.Module):
|
||||
"""Tensor-parallel MLP mapping merged vision features to text width."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vision_hidden_size: int,
|
||||
text_hidden_size: int,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if (
|
||||
isinstance(quant_config, Fp8Config)
|
||||
and prefix in quant_config.ignored_layers
|
||||
):
|
||||
quant_config = None
|
||||
|
||||
self.linear_fc1 = ColumnParallelLinear(
|
||||
vision_hidden_size,
|
||||
text_hidden_size,
|
||||
bias=True,
|
||||
quant_config=quant_config,
|
||||
prefix=maybe_prefix(prefix, "linear_fc1"),
|
||||
)
|
||||
self.act_fn = nn.GELU()
|
||||
self.linear_fc2 = RowParallelLinear(
|
||||
text_hidden_size,
|
||||
text_hidden_size,
|
||||
bias=True,
|
||||
quant_config=quant_config,
|
||||
prefix=maybe_prefix(prefix, "linear_fc2"),
|
||||
)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
hidden_states, _ = self.linear_fc1(hidden_states)
|
||||
hidden_states = self.act_fn(hidden_states)
|
||||
hidden_states, _ = self.linear_fc2(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
@MULTIMODAL_REGISTRY.register_processor(
|
||||
BailingMoeV3VLMultiModalProcessor,
|
||||
info=BailingMoeV3VLProcessingInfo,
|
||||
dummy_inputs=BailingMoeV3VLDummyInputsBuilder,
|
||||
)
|
||||
class BailingMoeV3VLForConditionalGeneration(
|
||||
nn.Module,
|
||||
HasInnerState,
|
||||
IsHybrid,
|
||||
SupportsMultiModal,
|
||||
SupportsPP,
|
||||
SupportsMRoPE,
|
||||
):
|
||||
"""Native vLLM wrapper for ``BailingMoeV3VLConfig`` checkpoints."""
|
||||
|
||||
packed_modules_mapping = {
|
||||
**BailingMoeV3ForCausalLM.packed_modules_mapping,
|
||||
# The vision checkpoint already stores Q/K/V in one tensor.
|
||||
"qkv": ["qkv"],
|
||||
}
|
||||
|
||||
hf_to_vllm_mapper = BailingMoeV3ForCausalLM.hf_to_vllm_mapper | WeightsMapper(
|
||||
orig_to_new_prefix={
|
||||
# Match module names as well as parameter names for quantization.
|
||||
"model.visual": "visual",
|
||||
"model.linear_proj": "linear_proj",
|
||||
"linear_proj.0": "linear_proj.linear_fc1",
|
||||
"linear_proj.2": "linear_proj.linear_fc2",
|
||||
"lm_head": "language_model.lm_head",
|
||||
"model.": "language_model.model.",
|
||||
}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
|
||||
if modality.startswith("image"):
|
||||
return IMAGE_PLACEHOLDER
|
||||
raise ValueError("Only image modality is supported")
|
||||
|
||||
@classmethod
|
||||
def get_mamba_state_shape_from_config(
|
||||
cls, vllm_config: VllmConfig
|
||||
) -> tuple[tuple[int, int], tuple[int, int, int]]:
|
||||
text_config = vllm_config.model_config.hf_config.text_config
|
||||
return BailingMoeV3ForCausalLM.get_mamba_state_shape_from_config(
|
||||
vllm_config.with_hf_config(text_config)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_mamba_state_dtype_from_config(
|
||||
cls, vllm_config: VllmConfig
|
||||
) -> tuple[torch.dtype, torch.dtype]:
|
||||
text_config = vllm_config.model_config.hf_config.text_config
|
||||
return BailingMoeV3ForCausalLM.get_mamba_state_dtype_from_config(
|
||||
vllm_config.with_hf_config(text_config)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_mamba_state_copy_func(cls) -> tuple:
|
||||
return BailingMoeV3ForCausalLM.get_mamba_state_copy_func()
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model") -> None:
|
||||
super().__init__()
|
||||
config = vllm_config.model_config.hf_config
|
||||
vision_config = config.vision_config
|
||||
text_config = config.text_config
|
||||
quant_config = vllm_config.quant_config
|
||||
|
||||
self.config = config
|
||||
self.model_config = vllm_config.model_config
|
||||
|
||||
merged_vision_size = vision_config.hidden_size * (
|
||||
vision_config.spatial_merge_size**2
|
||||
)
|
||||
with self._mark_tower_model(vllm_config, "image"):
|
||||
self.visual = BailingMoeV3VisionTransformer(
|
||||
vision_config,
|
||||
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
|
||||
quant_config=quant_config,
|
||||
prefix=maybe_prefix(prefix, "visual"),
|
||||
)
|
||||
self.linear_proj = BailingMoeV3VLProjector(
|
||||
vision_hidden_size=merged_vision_size,
|
||||
text_hidden_size=text_config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
prefix=maybe_prefix(prefix, "linear_proj"),
|
||||
)
|
||||
|
||||
with self._mark_language_model(vllm_config):
|
||||
self.language_model = BailingMoeV3ForCausalLM(
|
||||
vllm_config=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
|
||||
)
|
||||
|
||||
def _parse_and_validate_image_input(
|
||||
self, **kwargs: object
|
||||
) -> Qwen2_5_VLImageInputs | None:
|
||||
pixel_values = kwargs.pop("pixel_values", None)
|
||||
image_embeds = kwargs.pop("image_embeds", None)
|
||||
image_grid_thw = kwargs.pop("image_grid_thw", None)
|
||||
|
||||
if pixel_values is None and image_embeds is None:
|
||||
return None
|
||||
if pixel_values is not None:
|
||||
return Qwen2_5_VLImagePixelInputs(
|
||||
type="pixel_values",
|
||||
pixel_values=pixel_values,
|
||||
image_grid_thw=image_grid_thw,
|
||||
)
|
||||
return Qwen2_5_VLImageEmbeddingInputs(
|
||||
type="image_embeds",
|
||||
image_embeds=image_embeds,
|
||||
image_grid_thw=image_grid_thw,
|
||||
)
|
||||
|
||||
def _process_image_input(
|
||||
self, image_input: Qwen2_5_VLImageInputs
|
||||
) -> tuple[torch.Tensor, ...]:
|
||||
image_grid_thw = image_input["image_grid_thw"]
|
||||
assert image_grid_thw.ndim == 2
|
||||
|
||||
if image_input["type"] == "image_embeds":
|
||||
image_embeds = image_input["image_embeds"]
|
||||
else:
|
||||
pixel_values = image_input["pixel_values"].type(self.visual.dtype)
|
||||
vision_features = self.visual(
|
||||
pixel_values,
|
||||
grid_thw=image_grid_thw,
|
||||
)
|
||||
image_embeds = self.linear_proj(vision_features)
|
||||
|
||||
image_embeds = image_embeds.to(dtype=self.visual.dtype)
|
||||
merge_unit = self.visual.spatial_merge_size**2
|
||||
sizes = (image_grid_thw.prod(-1) // merge_unit).tolist()
|
||||
return image_embeds.split(sizes)
|
||||
|
||||
def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
|
||||
image_input = self._parse_and_validate_image_input(**kwargs)
|
||||
if image_input is None:
|
||||
return []
|
||||
return self._process_image_input(image_input)
|
||||
|
||||
def embed_input_ids(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
multimodal_embeddings: MultiModalEmbeddings | None = None,
|
||||
*,
|
||||
is_multimodal: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
inputs_embeds = self._embed_text_input_ids(
|
||||
input_ids,
|
||||
self.language_model.embed_input_ids,
|
||||
is_multimodal=is_multimodal,
|
||||
)
|
||||
if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
|
||||
return inputs_embeds
|
||||
|
||||
return _merge_multimodal_embeddings(
|
||||
inputs_embeds=inputs_embeds,
|
||||
multimodal_embeddings=multimodal_embeddings,
|
||||
is_multimodal=_require_is_multimodal(is_multimodal),
|
||||
)
|
||||
|
||||
def get_mrope_input_positions(
|
||||
self,
|
||||
input_tokens: list[int],
|
||||
mm_features: list[MultiModalFeatureSpec],
|
||||
) -> tuple[torch.Tensor, int]:
|
||||
position_blocks: list[np.ndarray] = []
|
||||
consumed = 0
|
||||
merge_size = self.config.vision_config.spatial_merge_size
|
||||
|
||||
for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):
|
||||
if mm_feature.modality != "image":
|
||||
raise ValueError(f"Unsupported modality: {mm_feature.modality}")
|
||||
|
||||
offset = mm_feature.mm_position.offset
|
||||
text_len = offset - consumed
|
||||
start = position_blocks[-1].max() + 1 if position_blocks else 0
|
||||
position_blocks.append(
|
||||
np.broadcast_to(np.arange(text_len), (3, text_len)) + start
|
||||
)
|
||||
|
||||
assert mm_feature.data is not None
|
||||
image_grid_thw = mm_feature.data["image_grid_thw"].data
|
||||
assert isinstance(image_grid_thw, torch.Tensor)
|
||||
t, h, w = image_grid_thw.tolist()
|
||||
assert t == 1, f"Image must have one temporal grid, got {t}"
|
||||
grid = np.indices((t, h // merge_size, w // merge_size))
|
||||
position_blocks.append(grid.reshape(3, -1) + text_len + start)
|
||||
consumed = offset + t * (h // merge_size) * (w // merge_size)
|
||||
|
||||
if consumed < len(input_tokens):
|
||||
start = position_blocks[-1].max() + 1 if position_blocks else 0
|
||||
text_len = len(input_tokens) - consumed
|
||||
position_blocks.append(
|
||||
np.broadcast_to(np.arange(text_len), (3, text_len)) + start
|
||||
)
|
||||
|
||||
if not position_blocks:
|
||||
positions = np.empty((3, 0), dtype=np.int64)
|
||||
return torch.from_numpy(positions), 0
|
||||
|
||||
positions = np.concatenate(position_blocks, axis=1).reshape(3, -1)
|
||||
position_delta = int(positions.max() + 1 - len(input_tokens))
|
||||
return torch.from_numpy(positions), position_delta
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**kwargs: object,
|
||||
) -> torch.Tensor | IntermediateTensors:
|
||||
if intermediate_tensors is not None:
|
||||
inputs_embeds = None
|
||||
return self.language_model(
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
intermediate_tensors=intermediate_tensors,
|
||||
inputs_embeds=inputs_embeds,
|
||||
)
|
||||
|
||||
def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
return self.language_model.compute_logits(hidden_states)
|
||||
|
||||
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)
|
||||
|
||||
def get_mm_mapping(self) -> MultiModelKeys:
|
||||
return MultiModelKeys.from_string_field(
|
||||
language_model="language_model",
|
||||
connector="linear_proj",
|
||||
tower_model="visual.",
|
||||
)
|
||||
|
||||
def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
|
||||
return num_image_tokens * self.visual.spatial_merge_size**2
|
||||
|
||||
def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
|
||||
return num_vision_tokens // self.visual.spatial_merge_size**2
|
||||
@@ -350,6 +350,10 @@ _MULTIMODAL_MODELS = {
|
||||
"AudioFlamingo3ForConditionalGeneration",
|
||||
),
|
||||
"BagelForConditionalGeneration": ("bagel", "BagelForConditionalGeneration"),
|
||||
"BailingMoeV3VLForConditionalGeneration": (
|
||||
"bailing_moe_v3_vl",
|
||||
"BailingMoeV3VLForConditionalGeneration",
|
||||
),
|
||||
"BeeForConditionalGeneration": ("bee", "BeeForConditionalGeneration"),
|
||||
"Blip2ForConditionalGeneration": ("blip2", "Blip2ForConditionalGeneration"),
|
||||
"Cohere2VisionForConditionalGeneration": (
|
||||
|
||||
@@ -71,10 +71,7 @@ class Ling3Parser(Glm47MoeParser):
|
||||
if not self.thinking_enabled:
|
||||
return None, model_output
|
||||
|
||||
reasoning, content = super().extract_reasoning(model_output, request)
|
||||
if reasoning and not content and "<tool_call>" not in model_output:
|
||||
return None, reasoning
|
||||
return reasoning, content
|
||||
return super().extract_reasoning(model_output, request)
|
||||
|
||||
|
||||
(
|
||||
|
||||
@@ -88,6 +88,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
|
||||
afmoe="AfmoeConfig",
|
||||
axk1="AXK1Config",
|
||||
bagel="BagelConfig",
|
||||
bailing_moe_v3_vl="BailingMoeV3VLConfig",
|
||||
chatglm="ChatGLMConfig",
|
||||
modernvbert="ColModernVBertConfig",
|
||||
colpali="ColPaliConfig",
|
||||
|
||||
@@ -18,6 +18,9 @@ _CLASS_TO_MODULE: dict[str, str] = {
|
||||
"AfmoeConfig": "vllm.transformers_utils.configs.afmoe",
|
||||
"AXK1Config": "vllm.transformers_utils.configs.AXK1",
|
||||
"BagelConfig": "vllm.transformers_utils.configs.bagel",
|
||||
"BailingMoeV3TextConfig": "vllm.transformers_utils.configs.bailing_moe_v3_vl",
|
||||
"BailingMoeV3VisionConfig": "vllm.transformers_utils.configs.bailing_moe_v3_vl",
|
||||
"BailingMoeV3VLConfig": "vllm.transformers_utils.configs.bailing_moe_v3_vl",
|
||||
"ChatGLMConfig": "vllm.transformers_utils.configs.chatglm",
|
||||
"ColModernVBertConfig": "vllm.transformers_utils.configs.colmodernvbert",
|
||||
"ColPaliConfig": "vllm.transformers_utils.configs.colpali",
|
||||
@@ -114,6 +117,9 @@ __all__ = [
|
||||
"AfmoeConfig",
|
||||
"AXK1Config",
|
||||
"BagelConfig",
|
||||
"BailingMoeV3TextConfig",
|
||||
"BailingMoeV3VisionConfig",
|
||||
"BailingMoeV3VLConfig",
|
||||
"ChatGLMConfig",
|
||||
"ColModernVBertConfig",
|
||||
"ColPaliConfig",
|
||||
|
||||
@@ -0,0 +1,345 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# Copyright 2025 Antgroup and The HuggingFace Inc. team. All rights reserved.
|
||||
# Adapted from
|
||||
# https://huggingface.co/inclusionAI/Ling-3.0-flash/blob/51fd444268df5267074ef0b289de0e1c4cb9b381/configuration_bailing_moe_v3.py
|
||||
"""Bailing V3 vision-language model configuration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
|
||||
|
||||
def _normalize_mrope_parameters(
|
||||
config: PretrainedConfig,
|
||||
mrope_section: list[int],
|
||||
) -> None:
|
||||
"""Add the top-level Bailing M-RoPE split to a text config."""
|
||||
rope_parameters = dict(getattr(config, "rope_parameters", None) or {})
|
||||
text_mrope_section = rope_parameters.get("mrope_section")
|
||||
if text_mrope_section is not None and list(text_mrope_section) != mrope_section:
|
||||
raise ValueError(
|
||||
"Bailing V3 VL has conflicting M-RoPE sections: "
|
||||
f"text={text_mrope_section}, top-level={mrope_section}"
|
||||
)
|
||||
|
||||
if "rope_type" not in rope_parameters and "type" not in rope_parameters:
|
||||
rope_parameters["rope_type"] = "default"
|
||||
rope_parameters.setdefault("rope_theta", getattr(config, "rope_theta", 10_000))
|
||||
rope_parameters["mrope_section"] = list(mrope_section)
|
||||
config.rope_parameters = rope_parameters
|
||||
|
||||
|
||||
def _build_layer_types(
|
||||
num_hidden_layers: int,
|
||||
layer_group_size: int,
|
||||
) -> list[str]:
|
||||
if layer_group_size <= 0:
|
||||
raise ValueError(
|
||||
f"Bailing V3 VL layer_group_size must be positive, got {layer_group_size}"
|
||||
)
|
||||
|
||||
grouped_layers = num_hidden_layers // layer_group_size * layer_group_size
|
||||
return [
|
||||
"linear_attention"
|
||||
if ((layer_idx + 1) % layer_group_size != 0 and layer_idx < grouped_layers)
|
||||
else "full_attention"
|
||||
for layer_idx in range(num_hidden_layers)
|
||||
]
|
||||
|
||||
|
||||
class BailingMoeV3TextConfig(PretrainedConfig):
|
||||
model_type = "bailing_hybrid"
|
||||
base_config_key = "text_config"
|
||||
keys_to_ignore_at_inference = ["past_key_values"]
|
||||
ignore_keys_at_rope_validation = {"mrope_section"}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
vocab_size: int = 157184,
|
||||
hidden_size: int = 2048,
|
||||
intermediate_size: int = 5120,
|
||||
num_hidden_layers: int = 20,
|
||||
num_attention_heads: int = 16,
|
||||
num_key_value_heads: int = 4,
|
||||
hidden_act: str = "silu",
|
||||
use_qkv_bias: bool = False,
|
||||
use_bias: bool = False,
|
||||
rms_norm_eps: float = 1e-6,
|
||||
tie_word_embeddings: bool = False,
|
||||
embedding_dropout: float = 0.0,
|
||||
attention_dropout: float = 0.0,
|
||||
output_dropout: float = 0.0,
|
||||
initializer_range: float = 0.02,
|
||||
max_position_embeddings: int = 32768,
|
||||
rope_theta: float = 600000.0,
|
||||
rope_parameters: dict[str, Any] | None = None,
|
||||
rope_scaling: dict[str, Any] | None = None,
|
||||
partial_rotary_factor: float | None = None,
|
||||
mrope_section: list[int] | tuple[int, ...] | None = None,
|
||||
use_cache: bool = True,
|
||||
max_window_layers: int = 20,
|
||||
pad_token_id: int = 156892,
|
||||
eos_token_id: int = 156892,
|
||||
num_experts: int = 256,
|
||||
num_shared_experts: int = 1,
|
||||
num_experts_per_tok: int = 8,
|
||||
n_group: int = 8,
|
||||
topk_group: int = 4,
|
||||
moe_intermediate_size: int = 512,
|
||||
moe_shared_expert_intermediate_size: int = 512,
|
||||
first_k_dense_replace: int = 1,
|
||||
head_dim: int | None = 128,
|
||||
output_router_logits: bool = False,
|
||||
use_qk_norm: bool = True,
|
||||
num_nextn_predict_layers: int = 0,
|
||||
mtp_loss_scaling_factor: float = 0,
|
||||
moe_router_enable_expert_bias: bool = True,
|
||||
routed_scaling_factor: float = 1.0,
|
||||
layer_group_size: int = 5,
|
||||
layer_types: list[str] | None = None,
|
||||
layers_block_type: list[str] | None = None,
|
||||
kv_lora_rank: int = 512,
|
||||
q_lora_rank: int | None = None,
|
||||
qk_rope_head_dim: int = 64,
|
||||
v_head_dim: int = 128,
|
||||
qk_nope_head_dim: int = 128,
|
||||
rope_interleave: bool = True,
|
||||
score_function: str = "sigmoid",
|
||||
scoring_func: str = "sigmoid",
|
||||
seq_aux: bool = True,
|
||||
topk_method: str = "noaux_tc",
|
||||
router_dtype: str = "fp32",
|
||||
gated_attention_proj_granularity_type: str | None = None,
|
||||
no_kda_lora: bool = False,
|
||||
kda_safe_gate: bool = False,
|
||||
kda_lower_bound: float | None = None,
|
||||
short_conv_kernel_size: int = 4,
|
||||
architectures: list[str] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
kwargs.pop("model_type", None)
|
||||
if architectures is None:
|
||||
architectures = ["BailingMoeV3ForCausalLM"]
|
||||
|
||||
if rope_parameters is not None:
|
||||
self.rope_parameters = dict(rope_parameters)
|
||||
elif rope_scaling is not None:
|
||||
self.rope_parameters = dict(rope_scaling)
|
||||
else:
|
||||
self.rope_parameters = {}
|
||||
if (
|
||||
"rope_type" not in self.rope_parameters
|
||||
and "type" not in self.rope_parameters
|
||||
):
|
||||
self.rope_parameters["rope_type"] = "default"
|
||||
self.rope_parameters.setdefault("rope_theta", rope_theta)
|
||||
if mrope_section is not None:
|
||||
_normalize_mrope_parameters(self, list(mrope_section))
|
||||
|
||||
self.vocab_size = vocab_size
|
||||
self.hidden_size = hidden_size
|
||||
self.intermediate_size = intermediate_size
|
||||
self.num_hidden_layers = num_hidden_layers
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.num_key_value_heads = num_key_value_heads
|
||||
self.hidden_act = hidden_act
|
||||
self.use_qkv_bias = use_qkv_bias
|
||||
self.use_bias = use_bias
|
||||
self.rms_norm_eps = rms_norm_eps
|
||||
self.embedding_dropout = embedding_dropout
|
||||
self.attention_dropout = attention_dropout
|
||||
self.output_dropout = output_dropout
|
||||
self.initializer_range = initializer_range
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.rope_theta = rope_theta
|
||||
self.partial_rotary_factor = partial_rotary_factor
|
||||
self.use_cache = use_cache
|
||||
self.max_window_layers = max_window_layers
|
||||
self.num_experts = num_experts
|
||||
self.num_shared_experts = num_shared_experts
|
||||
self.num_experts_per_tok = num_experts_per_tok
|
||||
self.n_group = n_group
|
||||
self.topk_group = topk_group
|
||||
self.moe_intermediate_size = moe_intermediate_size
|
||||
self.moe_shared_expert_intermediate_size = moe_shared_expert_intermediate_size
|
||||
self.first_k_dense_replace = first_k_dense_replace
|
||||
self.head_dim = head_dim or hidden_size // num_attention_heads
|
||||
self.output_router_logits = output_router_logits
|
||||
self.use_qk_norm = use_qk_norm
|
||||
self.num_nextn_predict_layers = num_nextn_predict_layers
|
||||
self.mtp_loss_scaling_factor = mtp_loss_scaling_factor
|
||||
self.moe_router_enable_expert_bias = moe_router_enable_expert_bias
|
||||
self.routed_scaling_factor = routed_scaling_factor
|
||||
self.layer_group_size = layer_group_size
|
||||
self.kv_lora_rank = kv_lora_rank
|
||||
self.q_lora_rank = q_lora_rank
|
||||
self.qk_rope_head_dim = qk_rope_head_dim
|
||||
self.v_head_dim = v_head_dim
|
||||
self.qk_nope_head_dim = qk_nope_head_dim
|
||||
self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
|
||||
self.rope_interleave = rope_interleave
|
||||
self.score_function = score_function
|
||||
self.scoring_func = scoring_func
|
||||
self.seq_aux = seq_aux
|
||||
self.topk_method = topk_method
|
||||
self.router_dtype = router_dtype
|
||||
self.gated_attention_proj_granularity_type = (
|
||||
gated_attention_proj_granularity_type
|
||||
)
|
||||
self.no_kda_lora = no_kda_lora
|
||||
self.kda_safe_gate = kda_safe_gate
|
||||
self.kda_lower_bound = kda_lower_bound
|
||||
self.short_conv_kernel_size = short_conv_kernel_size
|
||||
|
||||
expected_layer_types = _build_layer_types(
|
||||
num_hidden_layers,
|
||||
layer_group_size,
|
||||
)
|
||||
if layers_block_type is not None:
|
||||
legacy_type_mapping = {
|
||||
"attention": "full_attention",
|
||||
"mamba": "linear_attention",
|
||||
"conv": "linear_attention",
|
||||
}
|
||||
normalized_block_types = [
|
||||
legacy_type_mapping.get(layer_type, layer_type)
|
||||
for layer_type in layers_block_type
|
||||
]
|
||||
if layer_types is not None and layer_types != normalized_block_types:
|
||||
raise ValueError(
|
||||
"Bailing V3 VL has conflicting layer_types and layers_block_type"
|
||||
)
|
||||
layer_types = normalized_block_types
|
||||
if layer_types is not None and layer_types != expected_layer_types:
|
||||
raise ValueError(
|
||||
"Bailing V3 VL layer_types must match the layer_group_size schedule"
|
||||
)
|
||||
self.layer_types = expected_layer_types
|
||||
|
||||
super().__init__(
|
||||
architectures=architectures,
|
||||
pad_token_id=pad_token_id,
|
||||
eos_token_id=eos_token_id,
|
||||
tie_word_embeddings=tie_word_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
self.validate_layer_type()
|
||||
|
||||
|
||||
class BailingMoeV3VisionConfig(PretrainedConfig):
|
||||
model_type = "qwen3_moe_vit"
|
||||
base_config_key = "vision_config"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
depth: int = 27,
|
||||
hidden_size: int = 1152,
|
||||
hidden_act: str = "gelu_pytorch_tanh",
|
||||
intermediate_size: int = 4304,
|
||||
num_heads: int = 16,
|
||||
in_channels: int = 3,
|
||||
patch_size: int = 16,
|
||||
spatial_merge_size: int = 2,
|
||||
temporal_patch_size: int = 2,
|
||||
out_hidden_size: int = 3584,
|
||||
num_position_embeddings: int = 2304,
|
||||
deepstack_visual_indexes: list[int] | None = None,
|
||||
initializer_range: float = 0.02,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
kwargs.pop("model_type", None)
|
||||
super().__init__(**kwargs)
|
||||
self.depth = depth
|
||||
self.hidden_size = hidden_size
|
||||
self.hidden_act = hidden_act
|
||||
self.intermediate_size = intermediate_size
|
||||
self.num_heads = num_heads
|
||||
self.in_channels = in_channels
|
||||
self.patch_size = patch_size
|
||||
self.spatial_merge_size = spatial_merge_size
|
||||
self.temporal_patch_size = temporal_patch_size
|
||||
self.out_hidden_size = out_hidden_size
|
||||
self.num_position_embeddings = num_position_embeddings
|
||||
self.deepstack_visual_indexes = (
|
||||
[8, 16, 24]
|
||||
if deepstack_visual_indexes is None
|
||||
else deepstack_visual_indexes
|
||||
)
|
||||
self.initializer_range = initializer_range
|
||||
|
||||
|
||||
class BailingMoeV3VLConfig(PretrainedConfig):
|
||||
model_type = "bailing_moe_v3_vl"
|
||||
keys_to_ignore_at_inference = ["past_key_values"]
|
||||
sub_configs = {
|
||||
"text_config": BailingMoeV3TextConfig,
|
||||
"vision_config": BailingMoeV3VisionConfig,
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
text_config: BailingMoeV3TextConfig | dict[str, Any] | None = None,
|
||||
vision_config: BailingMoeV3VisionConfig | dict[str, Any] | None = None,
|
||||
image_token_id: int = 151655,
|
||||
video_token_id: int = 151656,
|
||||
vision_start_token_id: int = 151652,
|
||||
vision_end_token_id: int = 151653,
|
||||
tie_word_embeddings: bool = False,
|
||||
mrope_section: list[int] | tuple[int, ...] | None = None,
|
||||
architectures: list[str] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
kwargs.pop("model_type", None)
|
||||
if mrope_section is None:
|
||||
mrope_section = [12, 10, 10]
|
||||
self.mrope_section = list(mrope_section)
|
||||
|
||||
if text_config is None:
|
||||
text_config = BailingMoeV3TextConfig()
|
||||
elif isinstance(text_config, dict):
|
||||
text_config = dict(text_config)
|
||||
nested_mrope_section = text_config.pop("mrope_section", None)
|
||||
if (
|
||||
nested_mrope_section is not None
|
||||
and list(nested_mrope_section) != self.mrope_section
|
||||
):
|
||||
raise ValueError(
|
||||
"Bailing V3 VL has conflicting M-RoPE sections: "
|
||||
f"text={nested_mrope_section}, "
|
||||
f"top-level={self.mrope_section}"
|
||||
)
|
||||
text_config = BailingMoeV3TextConfig(**text_config)
|
||||
self.text_config = text_config
|
||||
_normalize_mrope_parameters(self.text_config, self.mrope_section)
|
||||
|
||||
if vision_config is None:
|
||||
vision_config = BailingMoeV3VisionConfig()
|
||||
elif isinstance(vision_config, dict):
|
||||
vision_config = BailingMoeV3VisionConfig(**vision_config)
|
||||
self.vision_config = vision_config
|
||||
|
||||
self.image_token_id = image_token_id
|
||||
self.video_token_id = video_token_id
|
||||
self.vision_start_token_id = vision_start_token_id
|
||||
self.vision_end_token_id = vision_end_token_id
|
||||
if architectures is None:
|
||||
architectures = ["BailingMoeV3VLForConditionalGeneration"]
|
||||
|
||||
super().__init__(
|
||||
architectures=architectures,
|
||||
tie_word_embeddings=tie_word_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BailingMoeV3TextConfig",
|
||||
"BailingMoeV3VisionConfig",
|
||||
"BailingMoeV3VLConfig",
|
||||
]
|
||||
@@ -1356,6 +1356,7 @@ class SpecDecodeBaseProposer:
|
||||
# handle multimodality
|
||||
assert hasattr(target_model, "config")
|
||||
if self.get_model_name(target_model) in [
|
||||
"BailingMoeV3VLForConditionalGeneration",
|
||||
"Cohere2VisionForConditionalGeneration",
|
||||
"Exaone4_5_ForConditionalGeneration",
|
||||
"GlmOcrForConditionalGeneration",
|
||||
|
||||
Reference in New Issue
Block a user