A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 21:00 UTC
[Feature] Add Qwen3-ForcedAligner support via token classification pooling (#35367)
Signed-off-by: haosdent <haosdent@gmail.com>
haosdent committed
Mar 29, 2026 at 08:27 UTC
d39b8daf5f6d220c68bb402806ad11e7be748883
7 files changed
+314
-2
docs/models/pooling_models/token_classify.md
+21
-2
@@ -15,7 +15,7 @@ Many classification models support both (sequence) classification and token clas
15
16
!!! note
17
18
- Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task (classify) is not
18
+ Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task (classify) is not
19
what you want, you need to manually specify it via `PoolerConfig(task="token_classify")` offline or
20
`--pooler-config.task token_classify` online.
21
@@ -29,6 +29,12 @@ Offline: [examples/pooling/token_classify/ner_offline.py](../../../examples/pool
29
30
Online: [examples/pooling/token_classify/ner_online.py](../../../examples/pooling/token_classify/ner_online.py)
31
32
+### Forced Alignment
33
+
34
+Forced alignment takes audio and reference text as input and produces word-level timestamps.
35
+
36
+Offline: [examples/pooling/token_classify/forced_alignment_offline.py](../../../examples/pooling/token_classify/forced_alignment_offline.py)
37
+
38
### Sparse retrieval (lexical matching)
39
40
The BAAI/bge-m3 model leverages token classification for sparse retrieval. For more information, see [this page](specific_models.md#baaibge-m3).
@@ -43,12 +49,25 @@ The BAAI/bge-m3 model leverages token classification for sparse retrieval. For m
49
| `Qwen3ForTokenClassification`<sup>C</sup> | Qwen3-based | `bd2lcco/Qwen3-0.6B-finetuned` | | |
50
| `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | \* | \* |
51
46
-<sup>C</sup> Automatically converted into a classification model via `--convert classify`. ([details](./README.md#model-conversion))
52
+<sup>C</sup> Automatically converted into a classification model via `--convert classify`. ([details](./README.md#model-conversion))
53
\* Feature support is the same as that of the original model.
54
55
If your model is not in the above list, we will try to automatically convert the model using
56
[as_seq_cls_model][vllm.model_executor.models.adapters.as_seq_cls_model]. By default, the class probabilities are extracted from the softmaxed hidden state corresponding to the last token.
57
58
+### Multimodal Models
59
+
60
+!!! note
61
+ For more information about multimodal models inputs, see [this page](../supported_models.md#list-of-multimodal-language-models).
62
+
63
+| Architecture | Models | Inputs | Example HF Models | [LoRA](../../features/lora.md) | [PP](../../serving/parallelism_scaling.md) |
64
+| --------------------------------------------- | ------------------- | ----------------- | ------------------------------------------ | ------------------------------ | ------------------------------------------ |
65
+| `Qwen3ASRForcedAlignerForTokenClassification` | Qwen3-ForcedAligner | T + A<sup>+</sup> | `Qwen/Qwen3-ForcedAligner-0.6B` (see note) | | ✅︎ |
66
+
67
+!!! note
68
+ Forced alignment usage requires `--hf-overrides '{"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]}'`.
69
+ Please refer to [examples/pooling/token_classify/forced_alignment_offline.py](../../../examples/pooling/token_classify/forced_alignment_offline.py).
70
+
71
### As Reward Models
72
73
Using token classification models as reward models. For details on reward models, see [Reward Models](reward.md).
examples/pooling/token_classify/forced_alignment_offline.py
new
+90
@@ -0,0 +1,90 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# Adapted from Qwen3-ForcedAligner inference:
4
+# https://github.com/QwenLM/Qwen3-ASR
5
+
6
+"""
7
+Offline forced alignment example using Qwen3-ForcedAligner-0.6B.
8
+
9
+Forced alignment takes audio and reference text as input and produces
10
+word-level timestamps. The model predicts a time bin at each <timestamp>
11
+token position; multiplying by ``timestamp_segment_time`` gives milliseconds.
12
+
13
+Usage::
14
+
15
+ python forced_alignment_offline.py \
16
+ --model Qwen/Qwen3-ForcedAligner-0.6B
17
+"""
18
+
19
+from argparse import Namespace
20
+
21
+import numpy as np
22
+
23
+from vllm import LLM, EngineArgs
24
+from vllm.utils.argparse_utils import FlexibleArgumentParser
25
+
26
+
27
+def parse_args():
28
+ parser = FlexibleArgumentParser()
29
+ parser = EngineArgs.add_cli_args(parser)
30
+ parser.set_defaults(
31
+ model="Qwen/Qwen3-ForcedAligner-0.6B",
32
+ runner="pooling",
33
+ enforce_eager=True,
34
+ hf_overrides={"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]},
35
+ )
36
+ return parser.parse_args()
37
+
38
+
39
+def build_prompt(words: list[str]) -> str:
40
+ """Build the forced alignment prompt from a word list.
41
+
42
+ Format: <|audio_start|><|audio_pad|><|audio_end|>
43
+ word1<timestamp><timestamp>word2<timestamp><timestamp>...
44
+ """
45
+ body = "<timestamp><timestamp>".join(words) + "<timestamp><timestamp>"
46
+ return f"<|audio_start|><|audio_pad|><|audio_end|>{body}"
47
+
48
+
49
+def main(args: Namespace):
50
+ llm = LLM(**vars(args))
51
+
52
+ config = llm.llm_engine.vllm_config.model_config.hf_config
53
+ timestamp_token_id = config.timestamp_token_id
54
+ timestamp_segment_time = config.timestamp_segment_time
55
+
56
+ # Example: align these words against a 5-second audio clip
57
+ words = ["Hello", "world"]
58
+ prompt = build_prompt(words)
59
+
60
+ # Use a 5-second silent audio as placeholder (replace with real audio)
61
+ sample_rate = 16000
62
+ audio = np.zeros(sample_rate * 5, dtype=np.float32)
63
+
64
+ outputs = llm.encode(
65
+ [{"prompt": prompt, "multi_modal_data": {"audio": audio}}],
66
+ pooling_task="token_classify",
67
+ )
68
+
69
+ for output in outputs:
70
+ logits = output.outputs.data # [num_tokens, classify_num]
71
+ predictions = logits.argmax(dim=-1)
72
+ token_ids = output.prompt_token_ids
73
+
74
+ # Extract timestamps at <timestamp> positions
75
+ ts_predictions = [
76
+ pred.item() * timestamp_segment_time
77
+ for tid, pred in zip(token_ids, predictions)
78
+ if tid == timestamp_token_id
79
+ ]
80
+
81
+ # Pair up start/end times per word
82
+ for i, word in enumerate(words):
83
+ start_ms = ts_predictions[i * 2]
84
+ end_ms = ts_predictions[i * 2 + 1]
85
+ print(f"{word:15s} {start_ms / 1000:.3f}s - {end_ms / 1000:.3f}s")
86
+
87
+
88
+if __name__ == "__main__":
89
+ args = parse_args()
90
+ main(args)
tests/models/multimodal/pooling/test_qwen3_asr_forced_aligner.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import numpy as np
5
+import pytest
6
+import torch
7
+
8
+MODEL = "Qwen/Qwen3-ForcedAligner-0.6B"
9
+CLASSIFY_NUM = 5000
10
+TIMESTAMP_TOKEN_ID = 151705
11
+
12
+
13
+def build_prompt(words: list[str]) -> str:
14
+ body = "<timestamp><timestamp>".join(words) + "<timestamp><timestamp>"
15
+ return f"<|audio_start|><|audio_pad|><|audio_end|>{body}"
16
+
17
+
18
+@pytest.mark.parametrize("model", [MODEL])
19
+@pytest.mark.parametrize("dtype", ["bfloat16"])
20
+@torch.inference_mode()
21
+def test_qwen3_forced_aligner(
22
+ vllm_runner,
23
+ model: str,
24
+ dtype: str,
25
+) -> None:
26
+ words = ["Hello", "world"]
27
+ prompt = build_prompt(words)
28
+
29
+ # 5-second silent audio at 16kHz
30
+ audio = np.zeros(16000 * 5, dtype=np.float32)
31
+
32
+ with vllm_runner(
33
+ model,
34
+ runner="pooling",
35
+ dtype=dtype,
36
+ enforce_eager=True,
37
+ max_model_len=512,
38
+ hf_overrides={
39
+ "architectures": [
40
+ "Qwen3ASRForcedAlignerForTokenClassification",
41
+ ],
42
+ },
43
+ ) as vllm_model:
44
+ outputs = vllm_model.llm.encode(
45
+ [{"prompt": prompt, "multi_modal_data": {"audio": audio}}],
46
+ pooling_task="token_classify",
47
+ )
48
+
49
+ # Validate output structure
50
+ assert len(outputs) == 1
51
+ logits = outputs[0].outputs.data
52
+ assert logits.dim() == 2
53
+ assert logits.shape[1] == CLASSIFY_NUM
54
+
55
+ # Validate timestamp extraction
56
+ token_ids = outputs[0].prompt_token_ids
57
+ predictions = logits.argmax(dim=-1)
58
+ ts_indices = [i for i, t in enumerate(token_ids) if t == TIMESTAMP_TOKEN_ID]
59
+
60
+ # 2 words x 2 timestamps each (start + end) = 4
61
+ assert len(ts_indices) == 4
62
+
63
+ ts_preds = [predictions[i].item() for i in ts_indices]
64
+ assert all(p >= 0 for p in ts_preds)
65
+ # end >= start for each word
66
+ assert ts_preds[1] >= ts_preds[0] # Hello
67
+ assert ts_preds[3] >= ts_preds[2] # world
tests/models/registry.py
+6
@@ -1094,6 +1094,12 @@ _MULTIMODAL_EXAMPLE_MODELS = {
1094
min_transformers_version="4.57",
1095
hf_overrides={"architectures": ["Qwen3ASRRealtimeGeneration"]},
1096
),
1097
+ "Qwen3ASRForcedAlignerForTokenClassification": _HfExamplesInfo(
1098
+ "Qwen/Qwen3-ForcedAligner-0.6B",
1099
+ max_model_len=4096,
1100
+ min_transformers_version="4.57",
1101
+ hf_overrides={"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]},
1102
+ ),
1103
"RForConditionalGeneration": _HfExamplesInfo("YannQi/R-4B", trust_remote_code=True),
1104
"SkyworkR1VChatModel": _HfExamplesInfo(
1105
"Skywork/Skywork-R1V-38B", trust_remote_code=True
vllm/model_executor/models/qwen3_asr_forced_aligner.py
new
+120
@@ -0,0 +1,120 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inference-only Qwen3-ASR ForcedAligner model (token classification)."""
4
+
5
+from collections.abc import Iterable
6
+
7
+import torch
8
+import torch.nn as nn
9
+
10
+from vllm.config import VllmConfig
11
+from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_classify
12
+from vllm.model_executor.models.interfaces_base import default_pooling_type
13
+from vllm.model_executor.models.qwen3_asr import (
14
+ Qwen3ASRDummyInputsBuilder,
15
+ Qwen3ASRForConditionalGeneration,
16
+ Qwen3ASRMultiModalProcessor,
17
+ Qwen3ASRProcessingInfo,
18
+)
19
+from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
20
+from vllm.multimodal import MULTIMODAL_REGISTRY
21
+from vllm.sequence import IntermediateTensors
22
+
23
+
24
+@default_pooling_type(tok_pooling_type="ALL")
25
+@MULTIMODAL_REGISTRY.register_processor(
26
+ Qwen3ASRMultiModalProcessor,
27
+ info=Qwen3ASRProcessingInfo,
28
+ dummy_inputs=Qwen3ASRDummyInputsBuilder,
29
+)
30
+class Qwen3ASRForcedAlignerForTokenClassification(
31
+ Qwen3ASRForConditionalGeneration,
32
+):
33
+ """Qwen3-ASR Forced Aligner model for per-token timestamp classification.
34
+
35
+ This model shares the audio tower and language model backbone with
36
+ Qwen3-ASR, but replaces the LM head with a classification head that
37
+ predicts time bins at ``<timestamp>`` token positions.
38
+
39
+ Usage::
40
+
41
+ llm = LLM(
42
+ model="Qwen/Qwen3-ForcedAligner-0.6B",
43
+ runner="pooling",
44
+ hf_overrides={
45
+ "architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]
46
+ },
47
+ )
48
+ outputs = llm.encode(
49
+ [{"prompt": prompt, "multi_modal_data": {"audio": audio}}],
50
+ pooling_task="token_classify",
51
+ )
52
+ """
53
+
54
+ is_pooling_model = True
55
+
56
+ # Map thinker.lm_head -> classifier (not language_model.lm_head)
57
+ hf_to_vllm_mapper = WeightsMapper(
58
+ orig_to_new_prefix={
59
+ "thinker.lm_head.": "classifier.",
60
+ "thinker.model.": "language_model.model.",
61
+ "thinker.": "",
62
+ }
63
+ )
64
+
65
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
66
+ super().__init__(vllm_config=vllm_config, prefix=prefix)
67
+
68
+ config = vllm_config.model_config.hf_config
69
+ thinker_config = config.thinker_config
70
+
71
+ # Remove the unused generation head created by the base class;
72
+ # the forced aligner uses a classifier head instead.
73
+ self.language_model.lm_head = None
74
+ self.language_model.logits_processor = None
75
+
76
+ self.classify_num = thinker_config.classify_num
77
+
78
+ # Classification head replaces lm_head for time-bin prediction.
79
+ # Use model dtype (not head_dtype which defaults to float32 for
80
+ # pooling models) to match the hidden state dtype.
81
+ self.classifier = nn.Linear(
82
+ thinker_config.text_config.hidden_size,
83
+ self.classify_num,
84
+ bias=False,
85
+ dtype=vllm_config.model_config.dtype,
86
+ )
87
+
88
+ # Token-level pooler to split per-token logits per request
89
+ pooler_config = vllm_config.model_config.pooler_config
90
+ assert pooler_config is not None
91
+ self.pooler = pooler_for_token_classify(pooler_config)
92
+
93
+ def forward(
94
+ self,
95
+ input_ids: torch.Tensor,
96
+ positions: torch.Tensor,
97
+ intermediate_tensors: IntermediateTensors | None = None,
98
+ inputs_embeds: torch.Tensor | None = None,
99
+ **kwargs: object,
100
+ ) -> torch.Tensor:
101
+ if intermediate_tensors is not None:
102
+ inputs_embeds = None
103
+
104
+ # Run through language model backbone (transformer layers only)
105
+ hidden_states = self.language_model.model(
106
+ input_ids,
107
+ positions,
108
+ intermediate_tensors,
109
+ inputs_embeds=inputs_embeds,
110
+ )
111
+
112
+ # Apply classification head -> [num_tokens, classify_num]
113
+ return self.classifier(hidden_states)
114
+
115
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
116
+ loader = AutoWeightsLoader(
117
+ self,
118
+ skip_prefixes=["talker.", "code2wav."],
119
+ )
120
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/registry.py
+4
@@ -292,6 +292,10 @@ _TOKEN_CLASSIFICATION_MODELS = {
292
"modernbert",
293
"ModernBertForTokenClassification",
294
),
295
+ "Qwen3ASRForcedAlignerForTokenClassification": (
296
+ "qwen3_asr_forced_aligner",
297
+ "Qwen3ASRForcedAlignerForTokenClassification",
298
+ ),
299
}
300
301
_SEQUENCE_CLASSIFICATION_MODELS = {
vllm/transformers_utils/configs/qwen3_asr.py
+6
@@ -342,12 +342,14 @@ class Qwen3ASRThinkerConfig(PretrainedConfig):
342
audio_start_token_id=151647,
343
user_token_id=872,
344
initializer_range=0.02,
345
+ classify_num=None,
346
**kwargs,
347
):
348
super().__init__(**kwargs)
349
self.user_token_id = user_token_id
350
self.audio_start_token_id = audio_start_token_id
351
self.initializer_range = initializer_range
352
+ self.classify_num = classify_num
353
354
if isinstance(audio_config, dict):
355
audio_config = Qwen3ASRAudioEncoderConfig(**audio_config)
@@ -406,6 +408,8 @@ class Qwen3ASRConfig(PretrainedConfig):
408
self,
409
thinker_config=None,
410
support_languages=None,
411
+ timestamp_token_id=None,
412
+ timestamp_segment_time=None,
413
**kwargs,
414
):
415
if thinker_config is None:
@@ -416,6 +420,8 @@ class Qwen3ASRConfig(PretrainedConfig):
420
421
self.thinker_config = Qwen3ASRThinkerConfig(**thinker_config)
422
self.support_languages = support_languages
423
+ self.timestamp_token_id = timestamp_token_id
424
+ self.timestamp_segment_time = timestamp_segment_time
425
super().__init__(**kwargs)
426
427
def get_text_config(self, decoder=False) -> "PretrainedConfig":