[Model] Support NVIDIA-NemotronLabs-VoiceChat-11B offline speech-to-speech (#5842)

Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

tests/model_executor/models/test_nemotron_voicechat_registration.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""L1 registration + config-shim tests for NemotronVoiceChat (CPU, no weights)."""
+
+import subprocess
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
+from vllm_omni.config.stage_config import StageExecutionType
+from vllm_omni.model_executor.models.nemotron_voicechat.configuration_nemotron_voicechat import (
+    NemotronVoiceChatConfig,
+)
+from vllm_omni.model_executor.models.registry import _OMNI_MODELS
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_ARCHS = (
+    "NemotronVoiceChatThinkerForConditionalGeneration",
+    "NemotronVoiceChatTalkerForConditionalGeneration",
+    "NemotronVoiceChatCode2Wav",
+)
+
+
+def _minimal_nemo_dict() -> dict:
+    return {
+        "data": {"frame_length": 0.08, "source_sample_rate": 16000, "target_sample_rate": 22050},
+        "model": {
+            "inference_speaker_name": "Aria",
+            "stt": {
+                "model": {
+                    "perception": {"encoder": {}, "preprocessor": {}, "modality_adapter": {}},
+                    "pretrained_llm": "nvidia/NVIDIA-Nemotron-Nano-9B-v2",
+                    "use_function_head": True,
+                }
+            },
+            "speech_generation": {
+                "model": {
+                    "tts_config": {"backbone_config": {"hidden_size": 1152, "sliding_window": 7500}},
+                    "codec_config": {"num_quantizers": 31, "codebook_size": 1024},
+                }
+            },
+        },
+    }
+
+
+def test_pipeline_registered_with_three_stages() -> None:
+    pipeline = OMNI_PIPELINES["nemotron_voicechat"]
+    assert OMNI_PIPELINES["nemotron_labs_voicechat"] is pipeline
+    assert len(pipeline.stages) == 3
+    thinker, talker, code2wav = pipeline.stages
+    assert thinker.execution_type == StageExecutionType.LLM_AR
+    assert thinker.final_output and thinker.final_output_type == "text"
+    assert thinker.owns_tokenizer
+    assert talker.execution_type == StageExecutionType.LLM_AR
+    assert talker.model_arch == "NemotronVoiceChatTalkerForConditionalGeneration"
+    assert talker.input_sources == (0,)
+    assert code2wav.execution_type == StageExecutionType.LLM_GENERATION
+    assert code2wav.final_output and code2wav.final_output_type == "audio"
+    assert code2wav.input_sources == (1,)
+    assert pipeline.validate() == []
+
+
+def test_architectures_registered() -> None:
+    for arch in _ARCHS:
+        assert arch in _OMNI_MODELS, f"{arch} missing from _OMNI_MODELS"
+
+
+def test_full_payload_whitelists_cover_code2wav_only() -> None:
+    from vllm_omni.core.sched.omni_scheduling_coordinator import _FULL_PAYLOAD_INPUT_STAGES
+
+    assert ("NemotronVoiceChatCode2Wav", "code2wav") in _FULL_PAYLOAD_INPUT_STAGES
+    # The thinker->talker hop is token-path only: the talker must NOT wait on a
+    # connector payload or stage 1 hangs.
+    assert ("NemotronVoiceChatTalkerForConditionalGeneration", "talker") not in _FULL_PAYLOAD_INPUT_STAGES
+
+
+def test_config_parses_nemo_layout_without_hub_access(monkeypatch) -> None:
+    # Provide the thinker text config explicitly so the test never touches HF.
+    cfg = NemotronVoiceChatConfig(
+        thinker_text_config={"model_type": "llama", "hidden_size": 4480, "num_hidden_layers": 56},
+        **_minimal_nemo_dict(),
+    )
+    assert cfg.get_text_config().hidden_size == 4480
+    assert cfg.talker_config.hidden_size == 1152
+    assert cfg.talker_config.sliding_window == 7500
+    assert cfg.code2wav_config.num_quantizers == 31
+    assert cfg.inference_speaker_name == "Aria"
+
+
+@pytest.mark.parametrize(
+    "mutation",
+    ["drop_stt", "drop_speech_generation", "drop_perception", "drop_codec"],
+)
+def test_config_rejects_non_voicechat_checkpoints(mutation: str) -> None:
+    raw = _minimal_nemo_dict()
+    if mutation == "drop_stt":
+        del raw["model"]["stt"]
+    elif mutation == "drop_speech_generation":
+        del raw["model"]["speech_generation"]
+    elif mutation == "drop_perception":
+        del raw["model"]["stt"]["model"]["perception"]
+    else:
+        del raw["model"]["speech_generation"]["model"]["codec_config"]
+    with pytest.raises(ValueError, match="NemotronVoiceChat|missing"):
+        NemotronVoiceChatConfig(
+            thinker_text_config={"model_type": "llama", "hidden_size": 4480},
+            **raw,
+        )
+
+
+def test_no_nemo_imports_in_nemotron_voicechat_tree() -> None:
+    """Dependency guard: the vendored tree must not import the nemo package.
+
+    Scoped to the nemotron_voicechat model tree (other models may have their
+    own optional nemo integrations).
+    """
+    repo_root = Path(__file__).resolve().parents[3]
+    targets = [
+        str(repo_root / "vllm_omni/model_executor/models/nemotron_voicechat"),
+        str(repo_root / "vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py"),
+    ]
+    result = subprocess.run(
+        ["grep", "-rnE", r"^\s*(import nemo\b|from nemo(\.|\s))", *targets, "--include=*.py"],
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode != 0, f"Found nemo imports in the nemotron_voicechat tree:\n{result.stdout}"

tests/model_executor/models/test_nemotron_voicechat_shapes.py

@@ -0,0 +1,227 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Golden shape test for the NemotronVoiceChat frame-locked timeline contract.
+
+Pins the exact prefill/off-by-one arithmetic against values verified with the
+NeMo reference run on sample_general.wav (fp32, greedy): perception produced
+196 acoustic frames for the 249,734-sample 16 kHz fixture (the NeMo dump's
+250,760 was the loader's padded batch width; perception consumed the true
+length); the system prompt tokenized to 56 ids ([bos] + 54 + [eos]); NeMo's
+timeline T = 252 = 56 + 196; and our fp32 pipeline matched the reference text
+channel 196/196 tokens.
+"""
+
+import pytest
+
+from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
+    compute_acoustic_frame_count,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Perception config subset matching the checkpoint (mel 10 ms hop, n_fft 512,
+# dw-striding 8x CAUSAL subsampling, conv kernel 3). causal_downsampling is
+# what the real checkpoint sets; omitting it here is exactly what let the
+# non-causal padding formula pass review round 0 (helper vs itself).
+_STT_CFG = {
+    "perception": {
+        "preprocessor": {"sample_rate": 16000, "window_stride": 0.01, "window_size": 0.025, "n_fft": 512},
+        "encoder": {
+            "subsampling": "dw_striding",
+            "subsampling_factor": 8,
+            "subsampling_conv_kernel_size": 3,
+            "causal_downsampling": True,
+        },
+    }
+}
+
+# (num_16k_samples, expected 12.5 Hz frames), all verified against the REAL
+# vendored AudioPerceptionModule built from the checkpoint config. The first
+# row is the sample_general.wav acceptance fixture (196 frames -> the NeMo
+# T=252 timeline). 71271 is the reviewer's 4.45 s repro that the non-causal
+# formula undersized by one; 250760 is the NeMo loader's padded width, which
+# genuinely yields 197 (not the fixture's 196).
+_GOLDEN = [
+    (249734, 196),
+    (250760, 197),
+    (71271, 57),
+    (16000, 14),
+    (16000 * 2, 26),
+    (16000 * 4, 51),
+    (16000 * 8, 101),
+]
+
+
+@pytest.mark.parametrize(("num_samples", "expected_frames"), _GOLDEN)
+def test_acoustic_frame_count_golden(num_samples: int, expected_frames: int) -> None:
+    assert compute_acoustic_frame_count(_STT_CFG, num_samples) == expected_frames
+
+
+def test_frame_count_matches_vendored_subsampling_geometry() -> None:
+    """Pin the helper against the real ConvSubsampling, not against itself.
+
+    Builds the vendored subsampling module with the checkpoint geometry (both
+    causal and non-causal) and checks the helper's padding arithmetic against
+    the module's actual ``_left_padding + _right_padding`` and its
+    ``calc_length`` output over a spread of mel lengths.
+    """
+    import torch
+
+    from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.asr.subsampling import (
+        ConvSubsampling,
+        calc_length,
+    )
+
+    for causal in (True, False):
+        sub = ConvSubsampling(
+            subsampling="dw_striding",
+            subsampling_factor=8,
+            feat_in=128,
+            feat_out=8,
+            conv_channels=8,
+            is_causal=causal,
+        )
+        expected_paddings = sub._left_padding + sub._right_padding
+        cfg = {
+            "perception": {
+                "preprocessor": _STT_CFG["perception"]["preprocessor"],
+                "encoder": {
+                    "subsampling": "dw_striding",
+                    "subsampling_factor": 8,
+                    "subsampling_conv_kernel_size": 3,
+                    "causal_downsampling": causal,
+                },
+            }
+        }
+        for num_samples in (16000, 71271, 249734, 250760):
+            # Recompute the helper's mel length, then ask the module's own
+            # calc_length what the subsampled length should be.
+            mel_len = (num_samples + 512 - 512) // 160 + 1
+            module_frames = int(
+                calc_length(
+                    torch.tensor([mel_len]),
+                    all_paddings=expected_paddings,
+                    kernel_size=3,
+                    stride=2,
+                    ceil_mode=False,
+                    repeat_num=3,
+                )[0]
+            )
+            assert compute_acoustic_frame_count(cfg, num_samples) == module_frames, (
+                f"helper diverged from vendored geometry (causal={causal}, samples={num_samples})"
+            )
+
+
+def test_timeline_max_model_len_validation() -> None:
+    from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
+        validate_timeline_fits_max_model_len,
+    )
+
+    # The acceptance fixture (56 + 1 + 196 = 253) fits the shipped 8192 limit.
+    validate_timeline_fits_max_model_len(56, 196, 8192)
+    # Exactly at the limit: the vLLM request occupies prompt + 1 (frame-0
+    # placeholder) + frames tokens, so 57 + 8135 == 8192 is allowed...
+    validate_timeline_fits_max_model_len(56, 8135, 8192)
+    # ...and one more frame must be rejected, naming every offending length.
+    with pytest.raises(ValueError) as excinfo:
+        validate_timeline_fits_max_model_len(56, 8136, 8192)
+    message = str(excinfo.value)
+    for fragment in (
+        "vllm_prefill_len=57",
+        "acoustic_frame_count=8136",
+        "vllm_total_len=8193",
+        "max_model_len=8192",
+        "logical_prompt_token_len=56",
+        "logical_total_len=8192",
+    ):
+        assert fragment in message
+
+
+def _bare_thinker():
+    import torch
+    from torch import nn
+
+    from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
+        NemotronVoiceChatThinkerForConditionalGeneration,
+    )
+
+    thinker = NemotronVoiceChatThinkerForConditionalGeneration.__new__(NemotronVoiceChatThinkerForConditionalGeneration)
+    nn.Module.__init__(thinker)
+    thinker._use_function_head = True
+    thinker._dtype = torch.float32
+    thinker.function_head = nn.Linear(8, 16, bias=False)
+    return thinker
+
+
+def test_postprocess_appends_function_token_to_timeline(monkeypatch: pytest.MonkeyPatch) -> None:
+    import torch
+
+    thinker = _bare_thinker()
+    hidden = torch.randn(3, 8)
+    expected = int(thinker.function_head(hidden[-1:]).argmax(dim=-1))
+
+    # Default: only the load-bearing prev-token; the cumulative debug timeline
+    # (a per-step host sync + growing concat) stays off.
+    update = thinker.postprocess(hidden)
+    assert int(update["nvc_prev_function_token"]) == expected
+    assert "nvc_function_tokens" not in update
+
+    monkeypatch.setenv("NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", "1")
+    # No existing timeline: starts one.
+    update = thinker.postprocess(hidden)
+    assert update["nvc_function_tokens"].tolist() == [expected]
+    # Existing timeline: appended, not replaced.
+    existing = torch.tensor([7, 8, 9], dtype=torch.long)
+    update = thinker.postprocess(hidden, nvc_function_tokens=existing)
+    assert update["nvc_function_tokens"].tolist() == [7, 8, 9, expected]
+
+
+def test_postprocess_skips_intermediate_prefill_chunks(monkeypatch: pytest.MonkeyPatch) -> None:
+    import torch
+
+    monkeypatch.setenv("NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", "1")
+    thinker = _bare_thinker()
+    hidden = torch.randn(4, 8)
+    # Intermediate chunk: 0 computed + 4 rows < prompt_len 10 -> no token.
+    assert thinker.postprocess(hidden, _omni_is_prefill=True, _omni_num_computed_tokens=0, _omni_prompt_len=10) == {}
+    # Final chunk reaches the last prefill position (acoustic frame 0).
+    update = thinker.postprocess(hidden, _omni_is_prefill=True, _omni_num_computed_tokens=6, _omni_prompt_len=10)
+    assert "nvc_prev_function_token" in update
+    assert update["nvc_function_tokens"].numel() == 1
+
+
+def test_prefill_contract_arithmetic() -> None:
+    # vllm_prefill_len = logical_prompt_token_len + 1 (the +1 is acoustic
+    # frame 0); decode steps == acoustic_frame_count; NeMo timeline
+    # T == prompt + frames. Values from the verified sample_general.wav run.
+    logical_prompt_len = 56
+    frames = compute_acoustic_frame_count(_STT_CFG, 249734)
+    vllm_prefill_len = logical_prompt_len + 1
+    assert vllm_prefill_len == 57
+    assert logical_prompt_len + frames == 252  # NeMo reference T
+    # The talker consumes the full timeline from t=1 and the producer trims
+    # the prompt region: rows P-1 onward of the (T-1)-row code stack.
+    talker_rows = (logical_prompt_len + frames) - 1
+    trimmed = talker_rows - (logical_prompt_len - 1)
+    assert trimmed == frames
+
+
+def test_torchaudio_mel_filterbank_matches_librosa() -> None:
+    """The perception frontend builds its mel filterbank with torchaudio.
+
+    torchaudio's melscale_fbanks(norm="slaney", mel_scale="slaney") implements
+    the same slaney-convention math as librosa.filters.mel (the NeMo reference
+    dependency) up to float32 rounding. The 1e-6 tolerance is ~4x the observed
+    2.6e-7 max diff; e2e this stays token- and WAV-identical to the reference
+    on the parity fixture (per-feature normalization absorbs it).
+    """
+    librosa = pytest.importorskip("librosa")
+    torchaudio = pytest.importorskip("torchaudio")
+
+    # The checkpoint's perception preprocessor config (16 kHz, n_fft=512, 128 mels).
+    ta = torchaudio.functional.melscale_fbanks(
+        n_freqs=257, f_min=0.0, f_max=8000.0, n_mels=128, sample_rate=16000, norm="slaney", mel_scale="slaney"
+    ).T.numpy()
+    reference = librosa.filters.mel(sr=16000, n_fft=512, n_mels=128, fmin=0.0, fmax=8000.0, norm="slaney")
+    assert ta.shape == reference.shape
+    assert abs(ta - reference).max() < 1e-6