anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Yuekai ZhangandClaude Fable 5 authored and GitHub committed 2026-08-12 16:25:48 +08:00
1 parent 9b0df3909a
commit 5be7541d01
46 files changed
+13034
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

examples/offline_inference/nemotron_voicechat/end2end.py
+229
View File
@@ -0,0 +1,229 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Offline speech-to-speech with NVIDIA-NemotronLabs-VoiceChat-11B.
Feeds a 16 kHz user utterance (plus a spoken-style system prompt) through the
3-stage nemotron_voicechat pipeline and writes the agent's reply as text and a
22.05 kHz WAV.
Usage:
python end2end.py \
--checkpoint /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
--wav /path/to/user_question_16k.wav \
--output-dir results/nemotron_voicechat
The system prompt defaults to the NeMo reference default. The text tokenizer
resolves from the ``nvidia/NVIDIA-Nemotron-Nano-9B-v2`` HF id (or the
``NEMOTRON_VOICECHAT_LLM_PATH`` env var for air-gapped runs).
Sizing the input WAV: the timeline is frame-locked, so the reply budget IS the
input duration — the acoustic channel trails the text channel, and a reply
that has not finished when the frames run out is truncated SILENTLY (no error;
ASR of the WAV will just disagree with the emitted text). Leave generous
trailing silence after the question (a question ending at ~4.5 s truncated in
an 8 s WAV but completed cleanly in a 16 s one). The script warns when the
text channel is still emitting non-PAD tokens near the last frame.
"""
from __future__ import annotations
import argparse
import json
import os
from pathlib import Path
import numpy as np
import soundfile as sf
# Warn if the text channel is still speaking within this many final frames
# (12 frames = ~1 s) — the frame-locked reply likely got cut off.
TRUNCATION_GUARD_FRAMES = 12
# The NeMo reference offline script's default system prompt.
DEFAULT_SYSTEM_PROMPT = (
"You are an AI voice assistant developed by NVIDIA. "
"Your name is NVIDIA Voice Chat. "
"Answer in a spoken, conversational style rather than a written one. "
"Do not repeat the same sentence over and over again. "
"Start the conversation by greeting the user."
)
def load_wav_16k_mono(path: str) -> np.ndarray:
wav, sr = sf.read(path, dtype="float32")
if wav.ndim > 1:
wav = wav.mean(axis=1)
if sr != 16000:
from vllm.multimodal.audio import resample_audio_scipy
wav = resample_audio_scipy(wav, orig_sr=sr, target_sr=16000)
return wav.astype(np.float32)
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--checkpoint", required=True, help="NVIDIA-NemotronLabs-VoiceChat-11B directory")
parser.add_argument("--wav", required=True, help="user speech (any sr; resampled to 16 kHz mono)")
parser.add_argument("--output-dir", default="results/nemotron_voicechat")
parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT)
parser.add_argument(
"--stage-configs-path",
default=None,
help=(
"Deploy yaml override. Pass vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml "
"for the async-chunk streaming pipeline (stages overlap; audio starts decoding "
"before the thinker finishes). Default: the bundled non-streaming yaml (bit-exact "
"parity path)."
),
)
args = parser.parse_args()
if not Path(args.wav).is_file():
raise FileNotFoundError(f"Input WAV not found: {args.wav}")
ckpt = Path(args.checkpoint)
raw_cfg = json.loads((ckpt / "config.json").read_text())
stt_cfg = raw_cfg["model"]["stt"]["model"]
from transformers import AutoTokenizer
from vllm.sampling_params import SamplingParams
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
compute_acoustic_frame_count,
)
tok_ref = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or stt_cfg.get(
"pretrained_llm", "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
)
tokenizer = AutoTokenizer.from_pretrained(tok_ref, trust_remote_code=False)
bos_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("bos_token", "<s>"))
eos_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("eos_token", "</s>"))
pad_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("pad_token", "<SPECIAL_12>"))
# NeMo encode_system_prompt: [bos] + tokens + [eos].
prompt_ids = [bos_id] + tokenizer.encode(args.system_prompt, add_special_tokens=False) + [eos_id]
wav = load_wav_16k_mono(args.wav)
n_frames = compute_acoustic_frame_count(stt_cfg, int(wav.shape[0]))
print(f"input: {wav.shape[0]} samples @16kHz -> {n_frames} acoustic frames; prompt {len(prompt_ids)} tokens")
inputs = {
# vLLM prompt = system prompt + one placeholder position (acoustic
# frame 0); the thinker replaces every position with fused embeddings.
"prompt_token_ids": prompt_ids + [pad_id],
"additional_information": {
# Plain float list: raw ndarrays do not survive the engine-core
# message serialization.
"nvc_audio": wav.astype(np.float32).tolist(),
"nvc_sr": 16000,
"nvc_prompt_token_ids": prompt_ids,
"nvc_expected_frames": n_frames,
},
}
sampling_params_list = [
# thinker: greedy, frame-locked generation length. NO min_tokens: the
# tokenizer's eos_token is <SPECIAL_12> — the frame-locked PAD/silence
# token — and vLLM's min_tokens masks "EOS" to -inf, which would forbid
# silence and force the agent to babble through the whole utterance.
SamplingParams(temperature=0.0, max_tokens=n_frames, ignore_eos=True, seed=0),
# talker: placeholder loop, stops on the stage's stop token. 16383:
# the 1-token placeholder prompt + max_tokens must fit the stage's
# max_model_len (16384).
SamplingParams(temperature=0.0, max_tokens=16383, stop_token_ids=[1], detokenize=False, seed=0),
# code2wav: single decode step.
SamplingParams(temperature=0.0, max_tokens=1, detokenize=False, seed=0),
]
omni_kwargs: dict = {}
if args.stage_configs_path:
omni_kwargs["stage_configs_path"] = args.stage_configs_path
omni = Omni(model=str(ckpt), trust_remote_code=True, **omni_kwargs)
outputs = omni.generate(inputs, sampling_params_list)
out_dir = Path(args.output_dir)
out_dir.mkdir(parents=True, exist_ok=True)
stem = Path(args.wav).stem
agent_text = None
audio = None
sr = 22050
for stage_outputs in outputs:
# One entry per final stage (text from the thinker, audio from code2wav);
# mirrors the qwen2_5_omni offline example's collection loop.
output = getattr(stage_outputs, "request_output", stage_outputs)
final_type = getattr(stage_outputs, "final_output_type", None)
completions = getattr(output, "outputs", None) or []
if not completions:
continue
completion = completions[0]
if final_type in (None, "text"):
text = getattr(completion, "text", None)
if text:
agent_text = text
token_ids = getattr(completion, "token_ids", None)
if token_ids:
(out_dir / f"{stem}_text_tokens.json").write_text(json.dumps(list(map(int, token_ids))))
# Frame lock: the reply budget IS the input length. If the text
# channel is still mid-sentence when the frames run out, the
# spoken answer is silently truncated — warn instead.
tail = list(map(int, token_ids))[-TRUNCATION_GUARD_FRAMES:]
if any(t != pad_id for t in tail):
print(
f"WARNING: the text channel emits non-PAD tokens within the last "
f"{TRUNCATION_GUARD_FRAMES} frames ({TRUNCATION_GUARD_FRAMES * 80} ms) — the reply "
"likely ran out of acoustic frames and is cut off mid-sentence. "
"Pad the input WAV with more trailing silence and rerun."
)
# Best-effort: dump the function-channel debug timeline if the
# engine surfaced the request info on this output.
info = getattr(output, "additional_information", None)
func_timeline = info.get("nvc_function_tokens") if isinstance(info, dict) else None
if func_timeline is not None:
values = func_timeline.tolist() if hasattr(func_timeline, "tolist") else list(func_timeline)
(out_dir / f"{stem}_function_tokens.json").write_text(json.dumps([int(v) for v in values]))
print(f"function-channel timeline: {len(values)} entries")
if final_type in (None, "audio"):
# multimodal_output is a MultimodalPayload (Mapping) — duck-type it.
mm = getattr(completion, "multimodal_output", None)
get = getattr(mm, "get", None)
if not callable(get):
continue
candidate = None
for key in ("audio", "model_outputs", "wav", "waveform"):
value = get(key)
if value is not None:
candidate = value
break
if isinstance(candidate, list):
candidate = candidate[0] if candidate else None
if candidate is not None:
arr = np.asarray(
candidate.float().cpu() if hasattr(candidate, "cpu") else candidate, dtype=np.float32
).reshape(-1)
if arr.size > 0:
audio = arr
sr_val = get("sr")
if sr_val is not None:
sr_item = sr_val[0] if isinstance(sr_val, list) else sr_val
sr = int(sr_item.item() if hasattr(sr_item, "item") else sr_item)
if agent_text is not None:
# The frame-locked text channel emits PAD tokens on silent frames;
# strip them (NeMo's post-inference detokenizer does the same).
pad_token = stt_cfg.get("pad_token", "<SPECIAL_12>")
agent_text = " ".join(agent_text.replace(pad_token, " ").split())
text_path = out_dir / f"{stem}_output.txt"
text_path.write_text(agent_text + "\n")
print(f"agent text -> {text_path}\n{agent_text}")
else:
print("WARNING: no agent text in outputs")
if audio is not None and audio.size > 0:
wav_path = out_dir / f"{stem}_output.wav"
sf.write(wav_path, audio, sr)
print(f"agent audio -> {wav_path} ({audio.size / sr:.2f}s @ {sr} Hz)")
else:
raise RuntimeError("No audio produced by the code2wav stage.")
if __name__ == "__main__":
main()
mkdocs.yml
+1
View File
@@ -111,6 +111,7 @@ plugins:
- "vllm_omni.entrypoints.openai" # avoid importing vllm in mkdocs building
- "vllm_omni.entrypoints.openai.protocol" # avoid importing vllm in mkdocs building
- "vllm_omni.experimental" # avoid importing optional serving deps (openai, fastapi, uvicorn, httpx) in mkdocs building
- "vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored" # vendored NeMo code; upstream docstrings trip griffe strict warnings
- mkdocstrings:
handlers:
python:
pyproject.toml
+13
View File
@@ -206,6 +206,12 @@ ignore = [
# These files keep upstream conventions (star imports, long lines, math-symbol
# argument names) so they stay easy to diff against upstream releases.
"vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py" = ["E402", "E501", "F403", "F405", "N803"]
# NemotronVoiceChat inference modules vendored from NVIDIA-NeMo/Speech (branch
# nemotron-labs-voicechat). Numerics-bearing code is kept byte-close to
# upstream to stay diffable and parity-safe, so upstream style (long lines,
# legacy typing, whitespace, unused-import shims) is exempted. TID251 and the
# F-correctness rules still apply.
"vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/**" = ["E501", "E721", "E741", "F401", "F841", "I001", "N803", "N814", "UP006", "UP007", "UP008", "UP011", "UP032", "UP035", "W291", "W293"]
"vllm_omni/experimental/fullduplex/joyvl/decision/prompts.py" = ["E501"] # System/summarizer prompts are long prose lines
[tool.ruff.lint.flake8-tidy-imports.banned-api]
@@ -314,6 +320,10 @@ extend-ignore-identifiers-re = [
".*[Nn][Dd]3.*",
".*updator.*",
"byt5.*",
# NemotronVoiceChat vendored triton kernel: dout/DOUT abbreviate d_out
# (output feature dim), not "doubt".
".*dout.*",
".*DOUT.*",
]
extend-ignore-re = [
"\\.\\.\\.n,d->\\.\\.\\.nd",
@@ -334,6 +344,9 @@ byt5 = "byt5"
# "seeked" is the standard HTML media element DOM event name (fires when a
# seek completes), used by the word-timestamps Gradio demo. Not a typo.
seeked = "seeked"
# Ryan Leary: author credited in the NemotronVoiceChat vendored NeMo
# features.py copyright header. A surname, not a typo of "leery".
Leary = "Leary"
# IndexTTS 2.5 language codes and speech-emotion-recognition event token.
fo = "fo"
ba = "ba"
recipes/NVIDIA/NemotronLabs-VoiceChat.md
+113
View File
@@ -0,0 +1,113 @@
# NemotronLabs VoiceChat 11B
> Offline speech-to-speech: a 16 kHz user utterance in, the agent's spoken reply
> (text + 22.05 kHz WAV) out, on a 3-stage vLLM-Omni pipeline.
## Summary
- Vendor: NVIDIA
- Model: [`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B)
- Task: Speech-to-speech voice chat. The model runs a frame-locked 12.5 Hz
timeline: a Conformer + NemotronH hybrid-Mamba thinker emits one text token
per acoustic frame, a Gemma3-1B EAR-TTS talker turns the text timeline into
31-quantizer RVQ code stacks, and an RVQ-VAE codec decodes them to audio.
- Mode: Offline single-turn inference (batch=1). Online/duplex serving,
function calling, and batch>1 are documented follow-ups.
- Maintainer: [`@yuekaizhang`](https://github.com/yuekaizhang)
## When to use this recipe
Use this recipe to run a single-turn voice-chat exchange with
`NVIDIA-NemotronLabs-VoiceChat-11B` on one GPU: you provide a user utterance as
a WAV file (any sample rate; it is resampled to 16 kHz mono) plus an optional
spoken-style system prompt, and get back the agent's reply as text and a
22.05 kHz WAV. The integration is NeMo-free at runtime — the perception
Conformer, EAR-TTS talker, and RVQ-VAE codec are vendored, dependency-stripped
NeMo modules (`nemo_vendored/`), so no `nemo_toolkit` install is needed.
## References
- Offline example:
[`examples/offline_inference/nemotron_voicechat/end2end.py`](../../examples/offline_inference/nemotron_voicechat/end2end.py)
- Model modules (thinker / talker / code2wav / vendored NeMo):
[`vllm_omni/model_executor/models/nemotron_voicechat/`](../../vllm_omni/model_executor/models/nemotron_voicechat/)
- Staged pipeline config:
[`vllm_omni/deploy/nemotron_labs_voicechat.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat.yaml)
- Upstream model card:
[`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B)
- Reference implementation: NVIDIA-NeMo/Speech, branch `nemotron-labs-voicechat`
## Pipeline
| stage | arch | dtype | role |
|---|---|---|---|
| 0 thinker | `NemotronVoiceChatThinkerForConditionalGeneration` (LLM_AR) | fp32 | WAV + system prompt -> frame-locked text-token timeline (+ function channel) |
| 1 talker | `NemotronVoiceChatTalker` (LLM_AR) | fp32 | text timeline -> 31-quantizer RVQ code stacks (one per 80 ms frame) |
| 2 code2wav | `NemotronVoiceChatCode2Wav` (LLM_GENERATION) | fp32 | RVQ-VAE decode -> 22.05 kHz PCM |
Stages 0/1 default to fp32 for exact parity with the NeMo reference
implementation (greedy decoding matches it token for token on the acceptance
fixture). The deploy yaml documents a ~2x-faster bf16 thinker option whose
output stayed within one word of the reference in testing.
## Hardware Support
## GPU
### 1x H100 80GB
#### Environment
- OS: Linux
- Python: 3.12
- vLLM version: 0.26.0
- vLLM-Omni version or commit: this PR / current `main`
#### Command
```bash
# Tokenizer: the checkpoint ships no HF tokenizer; it resolves from the
# nvidia/NVIDIA-Nemotron-Nano-9B-v2 HF id automatically. For air-gapped runs,
# point NEMOTRON_VOICECHAT_LLM_PATH at a local snapshot of that repo instead.
python examples/offline_inference/nemotron_voicechat/end2end.py \
--checkpoint /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
--wav /path/to/user_question.wav \
--output-dir results/nemotron_voicechat
```
#### Verification
```bash
ls results/nemotron_voicechat
# <stem>_output.txt the agent reply as text
# <stem>_output.wav the agent reply as 22.05 kHz audio
# <stem>_text_tokens.json the frame-locked text-token timeline
```
The reply text should read as a coherent spoken-style answer to the question in
the input WAV, and the WAV should transcribe to (approximately) the same text
with any ASR model.
#### Notes
- Memory usage: the shipped yaml runs all three stages on one GPU
(`gpu_memory_utilization` 0.62 / 0.12 / 0.06); peak usage is dominated by the
fp32 thinker. The fp32 default has a hard floor of roughly 43 GB of thinker
weights alone (9B backbone + 587M `embed_tokens` + 587M `function_head` +
0.6B Conformer), so 48 GB cards cannot run it — use the bf16 thinker option
documented in the deploy yaml on anything smaller than an 80 GB part.
- Input sizing: the timeline is frame-locked, so the reply budget IS the input
duration. The acoustic channel trails the text channel; if the WAV does not
carry enough trailing silence for the reply to finish, the spoken answer is
truncated silently. Leave generous trailing silence (a question ending at
~4.5 s truncated in an 8 s WAV but completed cleanly in 16 s); the offline
example warns when the text channel is still speaking near the last frame.
- Key flags: sampling is greedy end to end. The thinker is frame-locked —
`max_tokens` equals the acoustic frame count with `ignore_eos=True`. Do NOT
set `min_tokens` on the thinker: the tokenizer's EOS token is also the
frame-locked PAD/silence token, so masking it forces the model to babble
instead of pausing.
- The talker's `max_tokens` is 16383 (its stage prompt is one placeholder
token, and the stage context is 16384).
- Known limitations: batch=1 only; offline single turn (no online/duplex
serving yet); the function-call channel is decoded but not yet acted on.
recipes/README.md
+1
View File
@@ -66,6 +66,7 @@ recipes/
| [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
| [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
| [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB |
| [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB |
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
| [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
tests/config/test_omni_config.py
+1
View File
@@ -513,6 +513,7 @@ def test_sub_config_fields_match_structured_scopes():
"enable_prefix_caching",
"disable_hybrid_kv_cache_manager",
"mm_processor_cache_gb",
"mamba_ssm_cache_dtype",
}
assert {f.name for f in fields(OmniStageSchedulerConfig)} == {
"max_num_seqs",
tests/core/sched/test_omni_scheduling_coordinator.py
+3
View File
@@ -120,6 +120,9 @@ class TestFullPayloadCoordinatorSelection(unittest.TestCase):
("MiMoAudioModel", "code2wav"),
("Qwen3TTSCode2Wav", "code2wav"),
("MiniCPMO45Code2Wav", "code2wav"),
# nemotron_voicechat: sync-path RVQ-VAE decode waits for the
# talker's full [frames, 31] code-stack payload.
("NemotronVoiceChatCode2Wav", "code2wav"),
("CosyVoice3Model", "cosyvoice3_code2wav"),
# audex: sync-path TTS decoder (async_chunk: false deploys) and
# the TTA XCodec1 stage (always sync full-payload).
tests/model_executor/models/test_nemotron_voicechat_registration.py
+129
View File
@@ -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
+227
View File
@@ -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
tests/model_executor/stage_input_processors/test_nemotron_voicechat.py
+422
View File
@@ -0,0 +1,422 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import pytest
import torch
from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
_FULL_PAYLOAD_REPLACE_KEYS,
talker2code2wav_full_payload,
talker2code2wav_token_only,
thinker2talker_token_only,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
_PAD = 12
def _thinker_output(prompt_len: int, generated: list[int]):
return SimpleNamespace(
finished=True,
prompt_token_ids=list(range(100, 100 + prompt_len)) + [_PAD],
additional_information=None,
outputs=[SimpleNamespace(token_ids=generated, cumulative_token_ids=None)],
)
def test_thinker2talker_builds_pad_prompt_timeline() -> None:
generated = [7, 8, 9, 10]
inputs = thinker2talker_token_only([_thinker_output(prompt_len=5, generated=generated)])
assert len(inputs) == 1
prompt = inputs[0]
# 1-token placeholder prompt; the timeline rides additional_information.
assert prompt["prompt_token_ids"] == [0]
info = prompt["additional_information"]
assert info["nvc_logical_prompt_len"] == 5
assert info["nvc_frame_count"] == 4
assert info["nvc_text_timeline"] == [_PAD] * 5 + generated
def test_thinker2talker_skips_unfinished() -> None:
out = _thinker_output(3, [1])
out.finished = False
assert thinker2talker_token_only([out]) == []
def test_full_payload_trims_prompt_region_rows() -> None:
# Rows correspond to timeline steps t=1..T-1; with prompt_len P the NeMo
# trim keeps rows for t >= P, i.e. row index P-1 onward.
prompt_len = 4
frames, q = 10, 31
codes = torch.arange(frames * q, dtype=torch.long).reshape(frames, q) % 1024
request = SimpleNamespace(
additional_information={
"codes": {"audio": codes},
"nvc_logical_prompt_len": prompt_len,
},
additional_information_cpu=None,
)
payload = talker2code2wav_full_payload(request=request, is_finished=True)
assert payload.codes.audio.shape == (frames - (prompt_len - 1), q)
assert torch.equal(payload.codes.audio, codes[prompt_len - 1 :])
assert bool(payload.meta.finished)
def test_full_payload_without_codes_is_empty_finished() -> None:
# The sync producer must never emit a meta-only "pending" struct (it would
# be shipped as the real payload); with no codes it ships an empty finished
# payload regardless of the request's is_finished flag.
payload = talker2code2wav_full_payload(request=SimpleNamespace(), is_finished=False)
assert payload is not None
assert bool(payload.meta.finished)
assert payload.codes.audio.numel() == 0
def test_full_payload_replace_semantics_declared() -> None:
# Full-sequence tensors must REPLACE in the accumulator, not concatenate.
# The AR runner flattens nested payloads before accumulation, so the
# flattened "codes.audio" key is the one that matters.
assert "codes.audio" in _FULL_PAYLOAD_REPLACE_KEYS
assert "codes" in _FULL_PAYLOAD_REPLACE_KEYS
def test_full_payload_missing_prompt_len_fails_explicitly() -> None:
codes = torch.zeros((5, 31), dtype=torch.long)
request = SimpleNamespace(
additional_information={"codes": {"audio": codes}},
additional_information_cpu=None,
)
with pytest.raises(ValueError, match="nvc_logical_prompt_len"):
talker2code2wav_full_payload(request=request, is_finished=True)
def test_code2wav_distinguishes_empty_payload_from_missing_info() -> None:
from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_code2wav import (
NemotronVoiceChatCode2Wav,
)
extract = NemotronVoiceChatCode2Wav._codes_from_runtime_info
# Explicit empty payload (e.g. zero frames after the prompt trim): the
# empty tensor comes back as-is -> empty audio, never a placeholder-id
# decode.
empty = extract(None, {"codes": {"audio": torch.empty(0, dtype=torch.long)}})
assert isinstance(empty, torch.Tensor) and empty.numel() == 0
# Missing codes info entirely -> None (forward raises: placeholder
# input_ids are never decoded).
assert extract(None, {"meta": {}}) is None
assert extract(None, None) is None
# Normal payload passes through (nested and flattened key forms).
codes = torch.zeros((3, 31), dtype=torch.long)
assert extract(None, {"codes": {"audio": codes}}) is codes
assert extract(None, {"codes.audio": codes}) is codes
def test_validate_code_stack_negatives() -> None:
from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_code2wav import (
validate_code_stack,
)
good = torch.zeros((4, 31), dtype=torch.long)
assert validate_code_stack(good, 31, 1024).shape == (4, 31)
with pytest.raises(ValueError, match="not divisible"):
validate_code_stack(torch.zeros(100, dtype=torch.long), 31, 1024)
with pytest.raises(ValueError, match="shaped"):
validate_code_stack(torch.zeros((4, 30), dtype=torch.long), 31, 1024)
bad = good.clone()
bad[1, 5] = 1024
with pytest.raises(ValueError, match="out-of-range"):
validate_code_stack(bad, 31, 1024)
def test_tokenizer_resolution_failure_is_actionable(monkeypatch) -> None:
from transformers import AutoConfig
from vllm_omni.model_executor.models.nemotron_voicechat.configuration_nemotron_voicechat import (
NemotronVoiceChatConfig,
)
def _boom(*args, **kwargs):
raise OSError("offline")
monkeypatch.setattr(AutoConfig, "from_pretrained", _boom)
with pytest.raises(RuntimeError, match="NEMOTRON_VOICECHAT_LLM_PATH"):
NemotronVoiceChatConfig(
data={"frame_length": 0.08, "source_sample_rate": 16000, "target_sample_rate": 22050},
model={
"inference_speaker_name": "Aria",
"stt": {"model": {"perception": {"encoder": {}}, "pretrained_llm": "nvidia/nope"}},
"speech_generation": {"model": {"tts_config": {}, "codec_config": {}}},
},
)
def test_sanitize_tts_model_cfg_disables_pretrained_loads(monkeypatch) -> None:
from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_talker import (
sanitize_tts_model_cfg,
)
original = {
"pretrained_lm_name": "nvidia/NVIDIA-Nemotron-Nano-9B-v2",
"pretrained_codec_model": "/some/external/codec.ckpt",
"tts_config": {
"pretrained_text_name": "some/hub-llm",
"backbone_type": "gemma3_text",
"cas_config": {"pretrained_tokenizer_name": "some/hub-tokenizer", "keep": 1},
},
}
monkeypatch.delenv("NEMOTRON_VOICECHAT_LLM_PATH", raising=False)
cfg = sanitize_tts_model_cfg(original)
# External weight-load triggers are gone; construction is config-only.
assert "pretrained_text_name" not in cfg["tts_config"]
assert "pretrained_codec_model" not in cfg
assert "pretrained_tokenizer_name" not in cfg["tts_config"]["cas_config"]
assert cfg["tts_config"]["cas_config"]["keep"] == 1
# Tokenizer reference survives.
assert cfg["pretrained_lm_name"] == "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
# The shared stage config is never mutated (deep copy).
assert original["tts_config"]["pretrained_text_name"] == "some/hub-llm"
# Air-gapped override reroutes the tokenizer reference.
monkeypatch.setenv("NEMOTRON_VOICECHAT_LLM_PATH", "/local/nemotron")
cfg = sanitize_tts_model_cfg(original)
assert cfg["pretrained_lm_name"] == "/local/nemotron"
def test_code2wav_token_only_sizes_placeholder() -> None:
talker_output = SimpleNamespace(
finished=True,
outputs=[SimpleNamespace(token_ids=[0] * 9, cumulative_token_ids=None)],
)
inputs = talker2code2wav_token_only([talker_output])
assert len(inputs) == 1
assert len(inputs[0]["prompt_token_ids"]) == 9
# =========================
# Async-chunk (streaming) producers
# =========================
def _fake_transfer_manager(codec_chunk_frames: int | None = None):
extra = {}
if codec_chunk_frames is not None:
extra["codec_chunk_frames"] = codec_chunk_frames
return SimpleNamespace(
request_payload={},
connector=SimpleNamespace(config={"extra": extra}),
put_req_chunk={},
)
def _streaming_request(prompt_len: int, generated: list[int], info: dict | None = None):
return SimpleNamespace(
external_req_id="req-0",
prompt_token_ids=list(range(100, 100 + prompt_len)) + [_PAD],
output_token_ids=generated,
additional_information=info,
)
def test_thinker2talker_async_chunk_cumulative_timeline() -> None:
from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
thinker2talker_async_chunk,
)
tm = _fake_transfer_manager()
req = _streaming_request(prompt_len=3, generated=[7])
payload = thinker2talker_async_chunk(tm, {}, req, is_finished=False)
assert payload is not None
assert payload.ids.all == [_PAD] * 3 + [7]
assert payload.meta.num_processed_tokens == 3
assert not bool(payload.meta.finished)
# No new tokens since the last chunk -> skip.
assert thinker2talker_async_chunk(tm, {}, req, is_finished=False) is None
# More tokens -> longer cumulative timeline; final chunk sets finished.
req.output_token_ids = [7, 8, 9]
payload = thinker2talker_async_chunk(tm, {}, req, is_finished=True)
assert payload.ids.all == [_PAD] * 3 + [7, 8, 9]
assert bool(payload.meta.finished)
def test_talker2code2wav_async_chunk_cadence_and_left_context() -> None:
from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
talker2code2wav_async_chunk,
)
tm = _fake_transfer_manager(codec_chunk_frames=2)
prompt_len = 3
info = {"nvc_logical_prompt_len": prompt_len}
req = _streaming_request(prompt_len=prompt_len, generated=[], info=info)
def cumulative(rows: int) -> dict:
# Talker rows are timeline steps t=1..; prompt region occupies the
# first prompt_len-1 rows before acoustic frames start.
return {"codes": {"audio": torch.arange((prompt_len - 1 + rows) * 31).reshape(-1, 31)}}
# Only prompt-region rows so far -> nothing to ship.
assert talker2code2wav_async_chunk(tm, cumulative(0), req, is_finished=False) is None
# One new frame < chunk size 2 -> hold.
assert talker2code2wav_async_chunk(tm, cumulative(1), req, is_finished=False) is None
# Two frames -> first chunk, no left context.
payload = talker2code2wav_async_chunk(tm, cumulative(2), req, is_finished=False)
assert payload is not None
assert payload.meta.left_context_size == 0
assert payload.codes.audio.shape == (2, 31)
# Third frame + finished -> cumulative history with left context 2.
payload = talker2code2wav_async_chunk(tm, cumulative(3), req, is_finished=True)
assert payload.meta.left_context_size == 2
assert payload.codes.audio.shape == (3, 31)
assert bool(payload.meta.finished)
def test_talker2code2wav_async_chunk_requires_prompt_len() -> None:
from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
talker2code2wav_async_chunk,
)
tm = _fake_transfer_manager()
req = _streaming_request(prompt_len=3, generated=[], info=None)
with pytest.raises(ValueError, match="logical prompt length"):
talker2code2wav_async_chunk(tm, {"codes": {"audio": torch.zeros(5, 31, dtype=torch.long)}}, req)
# =========================
# Talker drain-per-wake (async scheduling review P1/P2)
# =========================
class _FakeTTSModelOutput:
past_key_values = None
class _FakeTTS:
"""Stub DuplexEARTTS: infer_codes_one_step returns a constant code row."""
def __init__(self, q: int = 31):
self.q = q
self.steps = 0
def set_init_inputs(self, speaker_name: str) -> None:
del speaker_name
def get_init_inputs(self, B: int) -> dict: # noqa: N803 — mirrors the vendored DuplexEARTTS signature
return {
"code": torch.zeros((B, 2, self.q), dtype=torch.long),
"subword_ids": torch.zeros((B, 3), dtype=torch.long),
}
def _get_generation_config(self, guidance_enabled: bool):
del guidance_enabled
return {}
def tts_model(self, **kwargs):
del kwargs
return _FakeTTSModelOutput()
def infer_codes_one_step(self, **kwargs):
del kwargs
self.steps += 1
return torch.full((1, 1, self.q), self.steps, dtype=torch.long), None
def _bare_talker():
from torch import nn
from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_talker import (
NemotronVoiceChatTalkerForConditionalGeneration,
)
talker = NemotronVoiceChatTalkerForConditionalGeneration.__new__(NemotronVoiceChatTalkerForConditionalGeneration)
nn.Module.__init__(talker)
talker.tts = _FakeTTS()
talker._sessions = {}
talker._hidden = 8
talker._dtype = torch.float32
talker._speaker_name = "Aria"
talker.config = SimpleNamespace(tts_cfg={"inference_guidance_enabled": False})
return talker
def _async_info(timeline: list[int], prompt_len: int, finished: bool, request_id: str = "req-0") -> dict:
return {
"request_id": request_id,
"additional_information": {
"ids": {"all": list(timeline)},
"meta": {
"num_processed_tokens": prompt_len,
"finished": torch.tensor(finished),
"request_id": request_id,
},
},
}
def test_talker_drains_all_received_positions_per_wake() -> None:
"""P2: chunk 0's PAD prompt region must not cost one upstream chunk per step,
and P1: a coalesced chunk (2+ new positions, 1 wake) must be fully drained."""
talker = _bare_talker()
pad, prompt_len = 12, 4
ids = torch.zeros(1, dtype=torch.long)
# Prefill wake: chunk 0 carries only the PAD prompt region.
info = _async_info([pad] * prompt_len, prompt_len, finished=False)
talker.preprocess(ids, None, _omni_is_prefill=True, **info)
session = talker._sessions["req-0"]
assert session["step"] == 1 and not session["sync_mode"]
# First decode wake: drains ALL prompt-region steps (t=1..P-1) at once.
_, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
assert session["step"] == prompt_len # t advanced 1 -> P in one wake
assert update["codes"]["audio"].shape == (prompt_len - 1, 31)
assert float(embeds[0, 0]) == 0.0 # not finished
# Coalesced chunk: TWO new frame tokens arrive in ONE wake (delayed save
# thread merged two thinker steps). Both must be consumed by this wake.
info2 = _async_info([pad] * prompt_len + [7, 8], prompt_len, finished=False)
_, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info2)
assert session["step"] == prompt_len + 2
assert update["codes"]["audio"].shape == (prompt_len + 1, 31)
assert float(embeds[0, 0]) == 0.0
# Zero-progress wake (no new positions, not finished): plain CONTINUE, no
# fabricated frames, no codes update.
_, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info2)
assert session["step"] == prompt_len + 2
assert update == {}
assert float(embeds[0, 0]) == 0.0
# Final chunk: one more token + finished marker -> drain + stop flag.
info3 = _async_info([pad] * prompt_len + [7, 8, 9], prompt_len, finished=True)
_, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info3)
assert session["step"] == prompt_len + 3
assert update["codes"]["audio"].shape == (prompt_len + 2, 31)
assert float(embeds[0, 0]) == 1.0 # stop
def test_talker_sync_mode_steps_once_per_wake() -> None:
"""The verified sync parity path keeps its one-TTS-step-per-engine-step cadence."""
talker = _bare_talker()
pad, prompt_len = 12, 3
ids = torch.zeros(1, dtype=torch.long)
info = {
"request_id": "req-1",
"additional_information": {
"nvc_text_timeline": [pad] * prompt_len + [7, 8],
"nvc_logical_prompt_len": prompt_len,
},
}
talker.preprocess(ids, None, _omni_is_prefill=True, **info)
session = talker._sessions["req-1"]
assert session["sync_mode"] and session["upstream_finished"]
for expected_step in (2, 3, 4):
_, embeds, _ = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
assert session["step"] == expected_step
# Last position -> stop flag on the final step.
_, embeds, _ = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
assert session["step"] == 5
assert float(embeds[0, 0]) == 1.0
vllm_omni/config/omni_config.py
+3
View File
@@ -128,6 +128,7 @@ class _CacheEngineOverrides(TypedDict, total=False):
enable_prefix_caching: bool
disable_hybrid_kv_cache_manager: bool
mm_processor_cache_gb: float
mamba_ssm_cache_dtype: str
class _SchedulerEngineOverrides(TypedDict, total=False):
@@ -377,6 +378,8 @@ class OmniStageCacheConfig:
enable_prefix_caching: bool | None = None
disable_hybrid_kv_cache_manager: bool | None = None
mm_processor_cache_gb: float | None = Field(default=None, ge=0.0)
# Hybrid-mamba SSM state dtype ("auto"/"float32"); vLLM CacheConfig field.
mamba_ssm_cache_dtype: str | None = None
@config
vllm_omni/config/pipeline_registry.py
+8
View File
@@ -88,6 +88,9 @@ from vllm_omni.model_executor.models.moss_tts.pipeline import (
MOSS_TTS_REALTIME_PIPELINE,
)
from vllm_omni.model_executor.models.moss_tts_nano.pipeline import MOSS_TTS_NANO_PIPELINE
from vllm_omni.model_executor.models.nemotron_voicechat.pipeline import (
NEMOTRON_VOICECHAT_PIPELINE,
)
from vllm_omni.model_executor.models.omnivoice.pipeline import OMNIVOICE_PIPELINE
from vllm_omni.model_executor.models.personaplex.pipeline import PERSONAPLEX_PIPELINE
from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import (
@@ -118,6 +121,11 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
"qwen2_5_omni": QWEN2_5_OMNI_PIPELINE,
"qwen2_5_omni_thinker_only": QWEN2_5_OMNI_THINKER_ONLY_PIPELINE,
"personaplex": PERSONAPLEX_PIPELINE,
"nemotron_voicechat": NEMOTRON_VOICECHAT_PIPELINE,
# Alias: lets bare `vllm-omni serve <NVIDIA-NemotronLabs-VoiceChat-11B dir>`
# auto-detect through the path-basename fallback (the checkpoint config.json
# has no model_type key).
"nemotron_labs_voicechat": NEMOTRON_VOICECHAT_PIPELINE,
"qwen3_omni_moe": resolve_qwen3_omni_pipeline,
"qwen3_tts": QWEN3_TTS_PIPELINE,
"step_audio_2": STEP_AUDIO2_PIPELINE,
vllm_omni/config/stage_config.py
+3
View File
@@ -357,6 +357,9 @@ class StageDeployConfig:
async_scheduling: bool | None = None
disable_hybrid_kv_cache_manager: bool | None = None
mm_processor_cache_gb: float | None = None
# Hybrid-mamba stages (e.g. the NemotronVoiceChat thinker's NemotronH
# backbone) pin the SSM state dtype; projected onto vLLM CacheConfig.
mamba_ssm_cache_dtype: str | None = None
# Generic compilation, profiling, tokenizer/config parsing, and model
# loading controls.
vllm_omni/core/sched/omni_scheduling_coordinator.py
+4
View File
@@ -51,6 +51,10 @@ _FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
("MiniCPMO45Code2Wav", "code2wav"),
# cosyvoice3: cosyvoice3_talker (Stage 0) -> cosyvoice3_code2wav (Stage 1).
("CosyVoice3Model", "cosyvoice3_code2wav"),
# nemotron_voicechat: talker (Stage 1) -> code2wav (Stage 2). Stage 2
# waits for the talker's full-payload code stacks; the thinker (Stage 0)
# -> talker (Stage 1) hop is token-path only and must NOT be listed.
("NemotronVoiceChatCode2Wav", "code2wav"),
# audex TTS sync path: thinker (Stage 0) -> streaming decoder (Stage 1).
# The default deploy is async_chunk; this covers async_chunk: false.
("AudexCode2Wav", "audex_code2wav"),
vllm_omni/deploy/nemotron_labs_voicechat.yaml
+81
View File
@@ -0,0 +1,81 @@
# NemotronVoiceChat-11B offline speech-to-speech deployment.
# 3 stages on one GPU: thinker (NemotronH AR) -> talker (EAR-TTS AR) -> code2wav
# (RVQ-VAE decode). thinker->talker is token-path only; talker->code2wav ships
# the [frames, 31] code stack via the shared-memory connector.
pipeline: nemotron_voicechat
async_chunk: false
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
stages:
- stage_id: 0 # thinker: speech -> frame-locked text (final text output)
max_num_seqs: 1
gpu_memory_utilization: 0.62
enforce_eager: true
enable_prefix_caching: false
async_scheduling: false
max_num_batched_tokens: 8192
max_model_len: 8192
# fp32 matches the NeMo reference exactly (196/196 token parity on the
# acceptance fixture). bf16 is ~2x faster and measured within one word of
# the reference on the same fixture — flip to bfloat16 if you prefer speed
# over exact parity.
dtype: float32
# NemotronH mamba SSM state in fp32 (matches NVIDIA's own vLLM deployment
# of this model; bf16 SSM state accumulates drift over long sequences).
mamba_ssm_cache_dtype: float32
devices: "0"
default_sampling_params:
# Greedy, frame-locked: the example sets max_tokens to the
# exact acoustic frame count per request; text EOS never stops generation.
temperature: 0.0
ignore_eos: true
seed: 0
- stage_id: 1 # talker: text timeline -> 31-quantizer code stacks
max_num_seqs: 1
gpu_memory_utilization: 0.12
enforce_eager: true
enable_prefix_caching: false
async_scheduling: false
skip_tokenizer_init: true
max_num_batched_tokens: 4096
max_model_len: 16384
# fp32 like the NeMo reference: the EAR-TTS MoG/codec-embedding path mixes
# fp32 buffers into the forward, and 1B params in fp32 are cheap.
dtype: float32
devices: "0"
output_connectors:
to_stage_2: connector_of_shared_memory
default_sampling_params:
# Placeholder-token loop; the stage emits stop token 1 after the last
# timeline position. 16383 (not 16384): the request also carries one
# placeholder prompt token, and prompt + max_tokens must fit
# max_model_len.
temperature: 0.0
max_tokens: 16383
stop_token_ids: [1]
detokenize: false
seed: 0
- stage_id: 2 # code2wav: codes -> 22.05 kHz PCM (fp32 codec)
max_num_seqs: 1
gpu_memory_utilization: 0.06
enforce_eager: true
enable_prefix_caching: false
async_scheduling: false
skip_tokenizer_init: true
max_num_batched_tokens: 65536
max_model_len: 65536
devices: "0"
input_connectors:
from_stage_1: connector_of_shared_memory
default_sampling_params:
temperature: 0.0
max_tokens: 1
detokenize: false
seed: 0
vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml
+84
View File
@@ -0,0 +1,84 @@
# NemotronVoiceChat-11B STREAMING speech-to-speech deployment (async-chunk).
#
# Same 3 stages as nemotron_labs_voicechat.yaml, but the stages overlap:
# the thinker streams its frame-locked text timeline to the talker per decode
# step (cumulative ids.all payloads), and the talker streams prompt-trimmed
# cumulative code stacks to code2wav every `codec_chunk_frames` frames. The
# codec re-decodes the full prefix per chunk and emits only the new samples
# (NeMo decode_one_audio_step semantics), so the concatenated stream is
# seam-free. First audio arrives after ~codec_chunk_frames thinker steps
# instead of after the whole pipeline completes.
#
# The non-streaming yaml remains the parity-acceptance path (bit-exact WAV);
# streamed audio is validated by ASR/WER instead (prefix decodes are not
# bit-identical to the one-shot decode of the final stack).
pipeline: nemotron_voicechat
async_chunk: true
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
# Streaming cadence for talker -> code2wav: emit a chunk once this many
# NEW frames (80 ms each) have accumulated; ~1 s of audio per chunk.
codec_chunk_frames: 13
stages:
- stage_id: 0 # thinker: speech -> frame-locked text (final text output)
max_num_seqs: 1
gpu_memory_utilization: 0.62
enforce_eager: true
enable_prefix_caching: false
async_scheduling: false
max_num_batched_tokens: 8192
max_model_len: 8192
dtype: float32
mamba_ssm_cache_dtype: float32
devices: "0"
output_connectors:
to_stage_1: connector_of_shared_memory
default_sampling_params:
temperature: 0.0
ignore_eos: true
seed: 0
- stage_id: 1 # talker: text timeline -> 31-quantizer code stacks
max_num_seqs: 1
gpu_memory_utilization: 0.12
enforce_eager: true
enable_prefix_caching: false
async_scheduling: false
skip_tokenizer_init: true
max_num_batched_tokens: 4096
max_model_len: 16384
dtype: float32
devices: "0"
input_connectors:
from_stage_0: connector_of_shared_memory
output_connectors:
to_stage_2: connector_of_shared_memory
default_sampling_params:
temperature: 0.0
max_tokens: 16383
stop_token_ids: [1]
detokenize: false
seed: 0
- stage_id: 2 # code2wav: codes -> 22.05 kHz PCM (fp32 codec)
max_num_seqs: 1
gpu_memory_utilization: 0.06
enforce_eager: true
enable_prefix_caching: false
async_scheduling: false
skip_tokenizer_init: true
max_num_batched_tokens: 65536
max_model_len: 65536
devices: "0"
input_connectors:
from_stage_1: connector_of_shared_memory
default_sampling_params:
temperature: 0.0
max_tokens: 1
detokenize: false
seed: 0
vllm_omni/engine/arg_utils.py
+22
View File
@@ -30,6 +30,11 @@ _ARCH_TO_MODEL_TYPE: dict[str, str] = {
# the HF config without manual overrides.
"PersonaPlexTalkerForConditionalGeneration": "personaplex",
"PersonaPlexCode2Wav": "personaplex",
# NemotronVoiceChat ships a NeMo-style config.json without model_type;
# create_model_config() patches model_type=nemotron_voicechat from the arch.
"NemotronVoiceChatThinkerForConditionalGeneration": "nemotron_voicechat",
"NemotronVoiceChatTalkerForConditionalGeneration": "nemotron_voicechat",
"NemotronVoiceChatCode2Wav": "nemotron_voicechat",
"VoxCPM2TalkerForConditionalGeneration": "voxcpm2",
}
@@ -56,6 +61,9 @@ def _register_omni_hf_configs() -> None:
MossTTSLocalConfig,
MossTTSRealtimeConfig,
)
from vllm_omni.model_executor.models.nemotron_voicechat.configuration_nemotron_voicechat import (
NemotronVoiceChatConfig,
)
from vllm_omni.model_executor.models.personaplex.configuration_personaplex import (
PersonaPlexConfig,
)
@@ -87,6 +95,7 @@ def _register_omni_hf_configs() -> None:
("moss_tts_realtime", MossTTSRealtimeConfig),
("qwen3_tts", Qwen3TTSConfig),
("personaplex", PersonaPlexConfig),
("nemotron_voicechat", NemotronVoiceChatConfig),
("cosyvoice3", CosyVoice3Config),
("glm_tts", GLMTTSConfig),
("omnivoice", OmniVoiceConfig),
@@ -324,6 +333,19 @@ class OmniEngineArgs(EngineArgs):
if model_type is not None:
self._patch_empty_hf_config(model_type)
# NemotronVoiceChat: the checkpoint's only tokenizer subfolder is the
# auxiliary rnnt_tokenizer/ (ASR loss vocab), which the generic
# subfolder auto-detect below would wrongly pick up. The text tokenizer
# is the Nemotron backbone's, resolved from a local override or its
# HF id (cache-friendly).
if not self.tokenizer and self.model_arch in (
"NemotronVoiceChatThinkerForConditionalGeneration",
"NemotronVoiceChatTalkerForConditionalGeneration",
"NemotronVoiceChatCode2Wav",
):
self.tokenizer = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
logger.info("NemotronVoiceChat: using text tokenizer from %s", self.tokenizer)
# Auto-detect tokenizer for models that store it in a subdirectory
# rather than the root (e.g. CosyVoice3 uses CosyVoice-BlankEN/).
if not self.tokenizer and self.model:
vllm_omni/model_executor/models/nemotron_voicechat/__init__.py
+19
View File
@@ -0,0 +1,19 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""NVIDIA NemotronLabs VoiceChat-11B: offline speech-to-speech 3-stage pipeline."""
from vllm_omni.model_executor.models.nemotron_voicechat.configuration_nemotron_voicechat import (
NemotronVoiceChatCode2WavConfig,
NemotronVoiceChatConfig,
NemotronVoiceChatTalkerConfig,
)
from vllm_omni.model_executor.models.nemotron_voicechat.pipeline import (
NEMOTRON_VOICECHAT_PIPELINE,
)
__all__ = [
"NemotronVoiceChatConfig",
"NemotronVoiceChatTalkerConfig",
"NemotronVoiceChatCode2WavConfig",
"NEMOTRON_VOICECHAT_PIPELINE",
]
vllm_omni/model_executor/models/nemotron_voicechat/configuration_nemotron_voicechat.py
+211
View File
@@ -0,0 +1,211 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""HF-compatible configs synthesized from the NeMo-style NemotronVoiceChat checkpoint.
The ``nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`` checkpoint ships a NeMo training
config (top-level ``data`` / ``model`` keys) instead of a transformers config:
there is no ``model_type`` or ``architectures``. vLLM-Omni's empty-config patching
(:meth:`OmniEngineArgs._patch_empty_hf_config`) injects
``model_type="nemotron_voicechat"`` and routes the raw dict here, where it is
parsed into per-stage sub-configs:
* ``get_text_config()`` -> the NemotronH backbone config (thinker stage; the AR
runner reads ``hidden_size`` etc. from it). Resolved from the checkpoint's
``model.stt.model.pretrained_llm`` HF id via the local HF cache, or from the
``NEMOTRON_VOICECHAT_LLM_PATH`` env var pointing at a local snapshot.
* ``talker_config`` -> Gemma3-based EAR-TTS backbone dims (talker stage).
* ``code2wav_config`` -> RVQ-VAE codec dims (code2wav stage).
The raw NeMo sections are preserved on the config (``stt_cfg`` / ``tts_cfg`` /
``nemo_data``) because the stage models consume them verbatim (perception layout,
fusion weights, TTS generation parameters, codec architecture).
"""
from __future__ import annotations
import os
from typing import Any
from transformers import AutoConfig, PretrainedConfig
_EXPECTED_STRUCTURE = (
"Expected a NemotronVoiceChat checkpoint config.json with the NeMo layout: "
"top-level 'model' containing both 'stt' (with 'model.stt.model.perception' "
"and 'model.stt.model.pretrained_llm') and 'speech_generation' (with "
"'model.speech_generation.model.tts_config' and '...codec_config')."
)
# Default HF id of the thinker text backbone; the checkpoint records the same id
# under model.stt.model.pretrained_llm.
_DEFAULT_LLM_ID = "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
# Env override for air-gapped runs: local directory containing the Nemotron
# backbone's config.json (and optionally tokenizer files).
_LLM_PATH_ENV = "NEMOTRON_VOICECHAT_LLM_PATH"
class NemotronVoiceChatTalkerConfig(PretrainedConfig):
"""Stage config for the EAR-TTS talker (28-layer Gemma3-style backbone)."""
model_type = "nemotron_voicechat_talker"
def __init__(self, **kwargs: Any) -> None:
backbone: dict[str, Any] = kwargs.pop("backbone_config", None) or {}
super().__init__(**kwargs)
self.backbone_config = backbone
self.hidden_size = int(backbone.get("hidden_size", 1152))
self.num_hidden_layers = int(backbone.get("num_hidden_layers", 28))
self.num_attention_heads = int(backbone.get("num_attention_heads", 16))
self.head_dim = int(backbone.get("head_dim", 72))
self.intermediate_size = int(backbone.get("intermediate_size", 4608))
self.sliding_window = int(backbone.get("sliding_window", 7500))
# The talker consumes/produces per-frame code stacks; vocab is nominal
# (codes are sampled by the vendored MOG head, not a vLLM sampler).
self.vocab_size = int(kwargs.get("vocab_size", 1024))
self.max_position_embeddings = int(backbone.get("max_position_embeddings", 32768))
def get_text_config(self, **_: Any) -> NemotronVoiceChatTalkerConfig:
return self
class NemotronVoiceChatCode2WavConfig(PretrainedConfig):
"""Stage config for the RVQ-VAE codec decoder."""
model_type = "nemotron_voicechat_code2wav"
def __init__(self, **kwargs: Any) -> None:
codec: dict[str, Any] = kwargs.pop("codec_config", None) or {}
super().__init__(**kwargs)
self.codec_config = codec
self.num_quantizers = int(codec.get("num_quantizers", 31))
self.codebook_size = int(codec.get("codebook_size", 1024))
self.latent_size = int(codec.get("latent_size", 512))
self.wav_to_token_ratio = int(codec.get("wav_to_token_ratio", 1764))
self.sample_rate = int(kwargs.get("sample_rate", 22050))
# The generation runner sizes buffers from hidden_size.
self.hidden_size = self.latent_size
self.num_attention_heads = 1
self.num_hidden_layers = 1
self.vocab_size = self.codebook_size
self.max_position_embeddings = int(kwargs.get("max_position_embeddings", 65536))
def get_text_config(self, **_: Any) -> NemotronVoiceChatCode2WavConfig:
return self
class NemotronVoiceChatConfig(PretrainedConfig):
"""Unified config: parses the raw NeMo dict into per-stage sub-configs."""
model_type = "nemotron_voicechat"
def __init__(self, **kwargs: Any) -> None:
nemo_model = kwargs.pop("model", None)
nemo_data = kwargs.pop("data", None)
thinker_text_config = kwargs.pop("thinker_text_config", None)
talker_config = kwargs.pop("talker_config", None)
code2wav_config = kwargs.pop("code2wav_config", None)
stt_cfg = kwargs.pop("stt_cfg", None)
tts_cfg = kwargs.pop("tts_cfg", None)
# Drop remaining NeMo training-only top-level sections.
for training_key in ("exp_manager", "trainer", "hf_export_dir", "name"):
kwargs.pop(training_key, None)
# transformers may call get_text_config() during super().__init__;
# give it a safe placeholder until the real sub-config is resolved.
self.thinker_text_config = None
super().__init__(**kwargs)
if nemo_model is not None:
stt = nemo_model.get("stt") or {}
speech_generation = nemo_model.get("speech_generation") or {}
if not (isinstance(stt, dict) and isinstance(stt.get("model"), dict)) or not (
isinstance(speech_generation, dict) and isinstance(speech_generation.get("model"), dict)
):
raise ValueError(
"config.json is missing the NemotronVoiceChat NeMo signature keys "
"'model.stt.model' and/or 'model.speech_generation.model'. " + _EXPECTED_STRUCTURE
)
stt_cfg = stt["model"]
tts_cfg = speech_generation["model"]
self.inference_speaker_name = nemo_model.get("inference_speaker_name") or "Aria"
self.tts_data = speech_generation.get("data") or {}
elif stt_cfg is None or tts_cfg is None:
# Bare construction (e.g. transformers plumbing); leave a skeleton
# that fails loudly if a stage model tries to use it.
self.inference_speaker_name = kwargs.get("inference_speaker_name", "Aria")
self.tts_data = {}
if stt_cfg is None or tts_cfg is None:
self.stt_cfg: dict[str, Any] = {}
self.tts_cfg: dict[str, Any] = {}
else:
if "perception" not in stt_cfg:
raise ValueError("'model.stt.model.perception' is missing from config.json. " + _EXPECTED_STRUCTURE)
if "tts_config" not in tts_cfg or "codec_config" not in tts_cfg:
raise ValueError(
"'model.speech_generation.model.tts_config'/'codec_config' are missing "
"from config.json. " + _EXPECTED_STRUCTURE
)
self.stt_cfg = dict(stt_cfg)
self.tts_cfg = dict(tts_cfg)
self.nemo_data = dict(nemo_data) if isinstance(nemo_data, dict) else {}
self.frame_length = float(self.nemo_data.get("frame_length", 0.08))
self.source_sample_rate = int(self.nemo_data.get("source_sample_rate", 16000))
self.target_sample_rate = int(self.nemo_data.get("target_sample_rate", 22050))
self.thinker_text_config = self._resolve_thinker_text_config(thinker_text_config)
self.talker_config = self._coerce(
talker_config,
NemotronVoiceChatTalkerConfig,
defaults={"backbone_config": (self.tts_cfg.get("tts_config") or {}).get("backbone_config") or {}},
)
self.code2wav_config = self._coerce(
code2wav_config,
NemotronVoiceChatCode2WavConfig,
defaults={
"codec_config": self.tts_cfg.get("codec_config") or {},
"sample_rate": self.target_sample_rate,
},
)
if not getattr(self, "architectures", None):
self.architectures = ["NemotronVoiceChatThinkerForConditionalGeneration"]
@staticmethod
def _coerce(value: Any, cls: type[PretrainedConfig], defaults: dict[str, Any]) -> PretrainedConfig:
if isinstance(value, cls):
return value
if isinstance(value, PretrainedConfig):
value = value.to_dict()
if isinstance(value, dict):
merged = dict(defaults)
merged.update(value)
return cls(**merged)
return cls(**defaults)
def _resolve_thinker_text_config(self, provided: Any) -> PretrainedConfig:
"""NemotronH backbone config: explicit > env-local > HF cache/hub."""
if isinstance(provided, PretrainedConfig):
return provided
if isinstance(provided, dict) and provided:
try:
return AutoConfig.for_model(**dict(provided))
except (KeyError, ValueError):
return PretrainedConfig(**provided)
if not self.stt_cfg:
# Skeleton config (bare construction); resolved lazily by the stage.
return PretrainedConfig()
llm_ref = os.environ.get(_LLM_PATH_ENV) or self.stt_cfg.get("pretrained_llm") or _DEFAULT_LLM_ID
try:
# nemotron_h is a native transformers model type; trust_remote_code=True
# would chase the auto_map .py modules and break air-gapped dirs that
# ship only config.json.
return AutoConfig.from_pretrained(llm_ref, trust_remote_code=False)
except Exception as exc:
raise RuntimeError(
f"Failed to load the NemotronVoiceChat thinker backbone config from {llm_ref!r}. "
f"Make sure the HF cache contains 'nvidia/NVIDIA-Nemotron-Nano-9B-v2' or set "
f"{_LLM_PATH_ENV} to a local directory with its config.json."
) from exc
def get_text_config(self, **_: Any) -> PretrainedConfig:
text_config = self.__dict__.get("thinker_text_config")
return text_config if text_config is not None else self
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/__init__.py
+34
View File
@@ -0,0 +1,34 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat); dependency-stripped
# copies of the NeMo inference modules needed to run NVIDIA-NemotronLabs-VoiceChat-11B.
from .duplex_ear_tts import DuplexEARTTS
from .ear_tts_commons import PreTrainedModel
from .ear_tts_vae_codec import RVQVAEModel
from .fusion import create_fusion_module
from .perception import AudioPerceptionModule, IdentityConnector
from .tokenizer import AutoTokenizer
__all__ = [
"AudioPerceptionModule",
"AutoTokenizer",
"DuplexEARTTS",
"IdentityConnector",
"PreTrainedModel",
"RVQVAEModel",
"create_fusion_module",
]
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/__init__.py
+25
View File
@@ -0,0 +1,25 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat); dependency-stripped
# copies of the NeMo ASR conformer modules needed by the perception encoder.
from .audio_preprocessing import AudioToMelSpectrogramPreprocessor
from .conformer_encoder import ConformerEncoder
__all__ = [
"AudioToMelSpectrogramPreprocessor",
"ConformerEncoder",
]
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/audio_preprocessing.py
+259
View File
@@ -0,0 +1,259 @@
# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/modules/audio_preprocessing.py; training-only code removed.
# Kept: AudioPreprocessor, AudioToMelSpectrogramPreprocessor (+ its config dataclass).
# Removed: AudioToMFCCPreprocessor, SpectrogramAugmentation, MaskedPatchAugmentation,
# CropOrPadSpectrogramAugmentation and numba/spec-augment imports (training-only).
import math
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Optional
import torch
from packaging import version
from ..compat import (
Exportable,
NeuralModule,
logging,
logging_mode,
typecheck,
)
from .features import FilterbankFeatures, FilterbankFeaturesTA
try:
import torchaudio
import torchaudio.functional
import torchaudio.transforms
TORCHAUDIO_VERSION = version.parse(torchaudio.__version__)
TORCHAUDIO_VERSION_MIN = version.parse("0.5")
HAVE_TORCHAUDIO = True
except ModuleNotFoundError:
HAVE_TORCHAUDIO = False
__all__ = [
"AudioToMelSpectrogramPreprocessor",
]
class AudioPreprocessor(NeuralModule, ABC):
"""
An interface for Neural Modules that performs audio pre-processing,
transforming the wav files to features.
"""
def __init__(self, win_length, hop_length):
super().__init__()
self.win_length = win_length
self.hop_length = hop_length
self.torch_windows = {
"hann": torch.hann_window,
"hamming": torch.hamming_window,
"blackman": torch.blackman_window,
"bartlett": torch.bartlett_window,
"ones": torch.ones,
None: torch.ones,
}
# Normally, when you call to(dtype) on a torch.nn.Module, all
# floating point parameters and buffers will change to that
# dtype, rather than being float32. The AudioPreprocessor
# classes, uniquely, don't actually have any parameters or
# buffers from what I see. In addition, we want the input to
# the preprocessor to be float32, but need to create the
# output in appropriate precision. We have this empty tensor
# here just to detect which dtype tensor this module should
# output at the end of execution.
self.register_buffer("dtype_sentinel_tensor", torch.tensor((), dtype=torch.float32), persistent=False)
@typecheck()
@torch.no_grad()
def forward(self, input_signal, length):
if input_signal.dtype != torch.float32:
logging.warning(
f"AudioPreprocessor received an input signal of dtype {input_signal.dtype}, rather than torch.float32. In sweeps across multiple datasets, we have found that the preprocessor is not robust to low precision mathematics. As such, it runs in float32. Your input will be cast to float32, but this is not necessarily enough to recovery full accuracy. For example, simply casting input_signal from torch.float32 to torch.bfloat16, then back to torch.float32 before running AudioPreprocessor causes drops in absolute WER of up to 0.1%. torch.bfloat16 simply does not have enough mantissa bits to represent enough values in the range [-1.0,+1.0] correctly.",
mode=logging_mode.ONCE,
)
processed_signal, processed_length = self.get_features(input_signal.to(torch.float32), length)
processed_signal = processed_signal.to(self.dtype_sentinel_tensor.dtype)
return processed_signal, processed_length
@abstractmethod
def get_features(self, input_signal, length):
# Called by forward(). Subclasses should implement this.
pass
class AudioToMelSpectrogramPreprocessor(AudioPreprocessor, Exportable):
"""Featurizer module that converts wavs to mel spectrograms.
Args:
sample_rate (int): Sample rate of the input audio data.
Defaults to 16000
window_size (float): Size of window for fft in seconds
Defaults to 0.02
window_stride (float): Stride of window for fft in seconds
Defaults to 0.01
n_window_size (int): Size of window for fft in samples
Defaults to None. Use one of window_size or n_window_size.
n_window_stride (int): Stride of window for fft in samples
Defaults to None. Use one of window_stride or n_window_stride.
window (str): Windowing function for fft. can be one of ['hann',
'hamming', 'blackman', 'bartlett']
Defaults to "hann"
normalize (str): Can be one of ['per_feature', 'all_features']; all
other options disable feature normalization. 'all_features'
normalizes the entire spectrogram to be mean 0 with std 1.
'pre_features' normalizes per channel / freq instead.
Defaults to "per_feature"
n_fft (int): Length of FT window. If None, it uses the smallest power
of 2 that is larger than n_window_size.
Defaults to None
preemph (float): Amount of pre emphasis to add to audio. Can be
disabled by passing None.
Defaults to 0.97
features (int): Number of mel spectrogram freq bins to output.
Defaults to 64
lowfreq (int): Lower bound on mel basis in Hz.
Defaults to 0
highfreq (int): Lower bound on mel basis in Hz.
Defaults to None
log (bool): Log features.
Defaults to True
log_zero_guard_type(str): Need to avoid taking the log of zero. There
are two options: "add" or "clamp".
Defaults to "add".
log_zero_guard_value(float, or str): Add or clamp requires the number
to add with or clamp to. log_zero_guard_value can either be a float
or "tiny" or "eps". torch.finfo is used if "tiny" or "eps" is
passed.
Defaults to 2**-24.
dither (float): Amount of white-noise dithering.
Defaults to 1e-5
pad_to (int): Ensures that the output size of the time dimension is
a multiple of pad_to.
Defaults to 16
frame_splicing (int): Defaults to 1
exact_pad (bool): If True, sets stft center to False and adds padding, such that num_frames = audio_length
// hop_length. Defaults to False.
pad_value (float): The value that shorter mels are padded with.
Defaults to 0
mag_power (float): The power that the linear spectrogram is raised to
prior to multiplication with mel basis.
Defaults to 2 for a power spec
rng : Random number generator
nb_augmentation_prob (float) : Probability with which narrowband augmentation would be applied to
samples in the batch.
Defaults to 0.0
nb_max_freq (int) : Frequency above which all frequencies will be masked for narrowband augmentation.
Defaults to 4000
use_torchaudio: Whether to use the `torchaudio` implementation.
mel_norm: Normalization used for mel filterbank weights.
Defaults to 'slaney' (area normalization)
stft_exact_pad: Deprecated argument, kept for compatibility with older checkpoints.
stft_conv: Deprecated argument, kept for compatibility with older checkpoints.
"""
def save_to(self, save_path: str):
pass
@classmethod
def restore_from(cls, restore_path: str):
pass
def __init__(
self,
sample_rate=16000,
window_size=0.02,
window_stride=0.01,
n_window_size=None,
n_window_stride=None,
window="hann",
normalize="per_feature",
n_fft=None,
preemph=0.97,
features=64,
lowfreq=0,
highfreq=None,
log=True,
log_zero_guard_type="add",
log_zero_guard_value=2**-24,
dither=1e-5,
pad_to=16,
frame_splicing=1,
exact_pad=False,
pad_value=0,
mag_power=2.0,
rng=None,
nb_augmentation_prob=0.0,
nb_max_freq=4000,
use_torchaudio: bool = False,
mel_norm="slaney",
stft_exact_pad=False, # Deprecated arguments; kept for config compatibility
stft_conv=False, # Deprecated arguments; kept for config compatibility
):
super().__init__(n_window_size, n_window_stride)
self._sample_rate = sample_rate
if window_size and n_window_size:
raise ValueError(f"{self} received both window_size and n_window_size. Only one should be specified.")
if window_stride and n_window_stride:
raise ValueError(f"{self} received both window_stride and n_window_stride. Only one should be specified.")
if window_size:
n_window_size = int(window_size * self._sample_rate)
if window_stride:
n_window_stride = int(window_stride * self._sample_rate)
# Given the long and similar argument list, point to the class and instantiate it by reference
if not use_torchaudio:
featurizer_class = FilterbankFeatures
else:
featurizer_class = FilterbankFeaturesTA
self.featurizer = featurizer_class(
sample_rate=self._sample_rate,
n_window_size=n_window_size,
n_window_stride=n_window_stride,
window=window,
normalize=normalize,
n_fft=n_fft,
preemph=preemph,
nfilt=features,
lowfreq=lowfreq,
highfreq=highfreq,
log=log,
log_zero_guard_type=log_zero_guard_type,
log_zero_guard_value=log_zero_guard_value,
dither=dither,
pad_to=pad_to,
frame_splicing=frame_splicing,
exact_pad=exact_pad,
pad_value=pad_value,
mag_power=mag_power,
rng=rng,
nb_augmentation_prob=nb_augmentation_prob,
nb_max_freq=nb_max_freq,
mel_norm=mel_norm,
stft_exact_pad=stft_exact_pad, # Deprecated arguments; kept for config compatibility
stft_conv=stft_conv, # Deprecated arguments; kept for config compatibility
)
def get_features(self, input_signal, length):
return self.featurizer(input_signal, length)
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/causal_convs.py
+154
View File
@@ -0,0 +1,154 @@
# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/submodules/causal_convs.py; training-only code removed.
from typing import Union
import torch
import torch.nn.functional as F
from torch import nn
__all__ = ["CausalConv2D", "CausalConv1D"]
class CausalConv2D(nn.Conv2d):
"""
A causal version of nn.Conv2d where each location in the 2D matrix would have no access to locations on its right or down
All arguments are the same as nn.Conv2d except padding which should be set as None
"""
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: int,
stride: int = 1,
padding: Union[str, int] = 0,
dilation: int = 1,
groups: int = 1,
bias: bool = True,
padding_mode: str = "zeros",
device=None,
dtype=None,
) -> None:
if padding is not None:
raise ValueError("Argument padding should be set to None for CausalConv2D.")
self._left_padding = kernel_size - 1
self._right_padding = stride - 1
padding = 0
super(CausalConv2D, self).__init__(
in_channels,
out_channels,
kernel_size,
stride,
padding,
dilation,
groups,
bias,
padding_mode,
device,
dtype,
)
def forward(
self,
x,
):
x = F.pad(x, pad=(self._left_padding, self._right_padding, self._left_padding, self._right_padding))
x = super().forward(x)
return x
class CausalConv1D(nn.Conv1d):
"""
A causal version of nn.Conv1d where each step would have limited access to locations on its right or left
All arguments are the same as nn.Conv1d except padding.
If padding is set None, then paddings are set automatically to make it a causal convolution where each location would not see any steps on its right.
If padding is set as a list (size of 2), then padding[0] would be used as left padding and padding[1] as right padding.
It would make it possible to control the number of steps to be accessible on the right and left.
This mode is not supported when stride > 1. padding[0]+padding[1] should be equal to (kernel_size - 1).
"""
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: int,
stride: int = 1,
padding: Union[str, int] = 0,
dilation: int = 1,
groups: int = 1,
bias: bool = True,
padding_mode: str = "zeros",
device=None,
dtype=None,
) -> None:
self.cache_drop_size = None
if padding is None:
self._left_padding = kernel_size - 1
self._right_padding = stride - 1
else:
if stride != 1 and padding != kernel_size - 1:
raise ValueError("No striding allowed for non-symmetric convolutions!")
if isinstance(padding, int):
self._left_padding = padding
self._right_padding = padding
elif isinstance(padding, list) and len(padding) == 2 and padding[0] + padding[1] == kernel_size - 1:
self._left_padding = padding[0]
self._right_padding = padding[1]
else:
raise ValueError(f"Invalid padding param: {padding}!")
self._max_cache_len = self._left_padding
super(CausalConv1D, self).__init__(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
padding=0,
dilation=dilation,
groups=groups,
bias=bias,
padding_mode=padding_mode,
device=device,
dtype=dtype,
)
def update_cache(self, x, cache=None):
if cache is None:
new_x = F.pad(x, pad=(self._left_padding, self._right_padding))
next_cache = cache
else:
new_x = F.pad(x, pad=(0, self._right_padding))
new_x = torch.cat([cache, new_x], dim=-1)
if self.cache_drop_size > 0:
next_cache = new_x[:, :, : -self.cache_drop_size]
else:
next_cache = new_x
next_cache = next_cache[:, :, -cache.size(-1) :]
return new_x, next_cache
def forward(self, x, cache=None):
x, cache = self.update_cache(x, cache=cache)
x = super().forward(x)
if cache is None:
return x
else:
return x, cache
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_encoder.py
+850
View File
@@ -0,0 +1,850 @@
# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/modules/conformer_encoder.py; training-only code removed.
import math
import random
from collections import OrderedDict
from dataclasses import dataclass
from typing import List, Optional, Set, Tuple
import torch
import torch.distributed
from omegaconf import DictConfig, ListConfig, open_dict
from torch import nn
from ..compat import (
AccessMixin,
CacheAwareStreamingConfig,
Exportable,
NeuralModule,
StreamingEncoder,
logging,
typecheck,
)
from .causal_convs import CausalConv1D
from .conformer_modules import ConformerLayer
from .multi_head_attention import (
LocalAttRelPositionalEncoding,
MultiHeadAttention,
PositionalEncoding,
RelPositionalEncoding,
RelPositionMultiHeadAttention,
RelPositionMultiHeadAttentionLongformer,
)
from .regularization_utils import compute_stochastic_depth_drop_probs
from .subsampling import (
ConvSubsampling,
StackingSubsampling,
SubsamplingReductionModule,
)
__all__ = ["ConformerEncoder"]
class ConformerEncoder(NeuralModule, StreamingEncoder, Exportable, AccessMixin):
"""
The encoder for ASR model of Conformer.
Based on this paper:
'Conformer: Convolution-augmented Transformer for Speech Recognition' by Anmol Gulati et al.
https://arxiv.org/abs/2005.08100
Args:
feat_in (int): the size of feature channels
n_layers (int): number of layers of ConformerBlock
d_model (int): the hidden size of the model
feat_out (int): the size of the output features
Defaults to -1 (means feat_out is d_model)
subsampling (str): the method of subsampling:
choices = ['vggnet', 'striding', 'dw-striding', 'stacking', 'stacking_norm']
Defaults to striding.
subsampling_factor (int): the subsampling factor which should be power of 2
Defaults to 4.
subsampling_conv_chunking_factor(int): optionally, force chunk inputs (helpful for large inputs)
Should be power of 2, 1 (auto-chunking, default), or -1 (no chunking)
subsampling_conv_channels (int): the size of the convolutions in the subsampling module
Defaults to -1 which would set it to d_model.
reduction (str, Optional): the method of reduction, choices=['pooling', 'striding']. If no value
is passed, then no reduction is performed and the models runs with the original 4x subsampling.
reduction_position (int, Optional): the index of the layer to apply reduction. If -1, apply reduction
at the end.
reduction_factor (int): the reduction factor which should be either 1 or a power of 2
Defaults to 1.
ff_expansion_factor (int): the expansion factor in feed forward layers
Defaults to 4.
self_attention_model (str): the type of the attention layer and positional encoding.
'rel_pos':
relative positional embedding and Transformer-XL
'rel_pos_local_attn':
relative positional embedding and Transformer-XL with local attention using
overlapping chunks. Attention context is determined by att_context_size parameter.
'abs_pos':
absolute positional embedding and Transformer
Default is rel_pos.
pos_emb_max_len (int): the maximum length of positional embeddings
Defaults to 5000
n_heads (int): number of heads in multi-headed attention layers
Defaults to 4.
att_context_size (List[Union[List[int],int]]): specifies the context sizes on each side.
Each context size should be a list of two integers like `[100, 100]`.
A list of context sizes like `[[100,100]`, `[100,50]]` can also be passed. -1 means unlimited context.
Defaults to `[-1, -1]`
att_context_probs (List[float]): a list of probabilities of each one of the att_context_size
when a list of them is passed. If not specified, uniform distribution is being used.
Defaults to None
att_context_style (str): 'regular' or 'chunked_limited'.
Defaults to 'regular'
xscaling (bool): enables scaling the inputs to the multi-headed attention layers by `sqrt(d_model)`.
Defaults to True.
untie_biases (bool): whether to not share (untie) the bias weights between layers of Transformer-XL
Defaults to True.
conv_kernel_size (int): the size of the convolutions in the convolutional modules
Defaults to 31.
conv_norm_type (str): the type of the normalization in the convolutional modules
Defaults to 'batch_norm'.
conv_context_size (list): it can be"causal" or a list of two integers
while `conv_context_size[0]+conv_context_size[1]+1==conv_kernel_size`.
`None` means `[(conv_kernel_size-1)//2`, `(conv_kernel_size-1)//2]`, and 'causal' means
`[(conv_kernel_size-1), 0]`.
Defaults to None.
conv_dual_mode (bool): specifies if convolution should be dual mode when dual_offline mode is being used.
When enables, the left half of the convolution kernel would get masked in streaming cases.
Defaults to False.
use_bias (bool): Use bias in all Linear and Conv1d layers from each ConformerLayer to improve
activation flow and stabilize training of huge models.
Defaults to True.
dropout (float): the dropout rate used in all layers except the attention layers
Defaults to 0.1.
dropout_pre_encoder (float): the dropout rate used before the encoder
Defaults to 0.1.
dropout_emb (float): the dropout rate used for the positional embeddings
Defaults to 0.1.
dropout_att (float): the dropout rate used for the attention layer
Defaults to 0.0.
stochastic_depth_drop_prob (float): if non-zero, will randomly drop
layers during training. The higher this value, the more often layers
are dropped. Defaults to 0.0.
stochastic_depth_mode (str): can be either "linear" or "uniform". If
set to "uniform", all layers have the same probability of drop. If
set to "linear", the drop probability grows linearly from 0 for the
first layer to the desired value for the final layer. Defaults to
"linear".
stochastic_depth_start_layer (int): starting layer for stochastic depth.
All layers before this will never be dropped. Note that drop
probability will be adjusted accordingly if mode is "linear" when
start layer is > 1. Defaults to 1.
global_tokens (int): number of tokens to be used for global attention.
Only relevant if self_attention_model is 'rel_pos_local_attn'.
Defaults to 0.
global_tokens_spacing (int): how far apart the global tokens are
Defaults to 1.
global_attn_separate (bool): whether the q, k, v layers used for global tokens should be separate.
Defaults to False.
use_pytorch_sdpa (bool): use torch sdpa instead of manual attention.
Defaults to False.
use_pytorch_sdpa_backends (list[str]): list of backend names to use in sdpa.
None or empty list means all backends. e.g. ["MATH"]
Defaults to None.
bypass_pre_encode: if True, skip the pre-encoder module and the `audio_signal` should be pre-encoded
embeddings. The `audio_signal` input supports two formats depending on the `bypass_pre_encode`
boolean flag. This determines the required format of the input variable `audio_signal`.
Defaults to `bypass_pre_encode=False`. `bypass_pre_encode=True` is used for the cases
where frame-level, context-independent embeddings are needed to be saved or reused.
(e.g., speaker cache in streaming speaker diarization)
sync_max_audio_length (bool): when true, performs NCCL all_reduce to allocate the same amount of memory for
positional encoding buffers on all GPUs. Disabling this setting may help with deadlocks in certain
scenarios such as model parallelism, or generally when this module is not being ran on some GPUs
as a part of the training step.
"""
def __init__(
self,
feat_in,
n_layers,
d_model,
feat_out=-1,
causal_downsampling=False,
subsampling="striding",
subsampling_factor=4,
subsampling_conv_chunking_factor=1,
subsampling_conv_channels=-1,
reduction=None,
reduction_position=None,
reduction_factor=1,
ff_expansion_factor=4,
self_attention_model="rel_pos",
n_heads=4,
att_context_size=None,
att_context_probs=None,
att_context_style="regular",
xscaling=True,
untie_biases=True,
pos_emb_max_len=5000,
conv_kernel_size=31,
conv_norm_type="batch_norm",
conv_context_size=None,
use_bias=True,
dropout=0.1,
dropout_pre_encoder=0.1,
dropout_emb=0.1,
dropout_att=0.0,
stochastic_depth_drop_prob: float = 0.0,
stochastic_depth_mode: str = "linear",
stochastic_depth_start_layer: int = 1,
global_tokens: int = 0,
global_tokens_spacing: int = 1,
global_attn_separate: bool = False,
use_pytorch_sdpa: bool = False,
use_pytorch_sdpa_backends=None,
sync_max_audio_length: bool = True,
):
super().__init__()
d_ff = d_model * ff_expansion_factor
self.d_model = d_model
self.n_layers = n_layers
self._feat_in = feat_in
self.att_context_style = att_context_style
self.subsampling_factor = subsampling_factor
self.subsampling_conv_chunking_factor = subsampling_conv_chunking_factor
self.self_attention_model = self_attention_model
self.global_tokens = global_tokens
self.global_attn_separate = global_attn_separate
self.global_tokens_spacing = global_tokens_spacing
self.use_pytorch_sdpa = use_pytorch_sdpa
if use_pytorch_sdpa_backends is None:
use_pytorch_sdpa_backends = []
self.use_pytorch_sdpa_backends = use_pytorch_sdpa_backends
self.sync_max_audio_length = sync_max_audio_length
# Setting up the att_context_size
(
self.att_context_size_all,
self.att_context_size,
self.att_context_probs,
self.conv_context_size,
) = self._calc_context_sizes(
att_context_style=att_context_style,
att_context_size=att_context_size,
att_context_probs=att_context_probs,
conv_context_size=conv_context_size,
conv_kernel_size=conv_kernel_size,
)
if xscaling:
self.xscale = math.sqrt(d_model)
else:
self.xscale = None
# Subsampling
if subsampling_conv_channels == -1:
subsampling_conv_channels = d_model
if subsampling and subsampling_factor > 1:
if subsampling in ["stacking", "stacking_norm"]:
# stacking_norm has an extra layer norm after stacking comparing to stacking
self.pre_encode = StackingSubsampling(
subsampling_factor=subsampling_factor,
feat_in=feat_in,
feat_out=d_model,
norm=True if subsampling == "stacking_norm" else False,
)
else:
self.pre_encode = ConvSubsampling(
subsampling=subsampling,
subsampling_factor=subsampling_factor,
feat_in=feat_in,
feat_out=d_model,
conv_channels=subsampling_conv_channels,
subsampling_conv_chunking_factor=subsampling_conv_chunking_factor,
activation=nn.ReLU(True),
is_causal=causal_downsampling,
)
else:
self.pre_encode = nn.Linear(feat_in, d_model)
# Reduction
if reduction and reduction_factor > 1:
assert reduction_position >= -1 and reduction_position < n_layers
self.reduction_subsampling = SubsamplingReductionModule(
reduction=reduction,
d_model=d_model,
reduction_factor=reduction_factor,
)
self.reduction_position = reduction_position
else:
self.reduction_subsampling = None
self.reduction_position = None
self._feat_out = d_model
# Biases for relative positional encoding
if not untie_biases and self_attention_model == "rel_pos":
d_head = d_model // n_heads
pos_bias_u = nn.Parameter(torch.Tensor(n_heads, d_head))
pos_bias_v = nn.Parameter(torch.Tensor(n_heads, d_head))
nn.init.zeros_(pos_bias_u)
nn.init.zeros_(pos_bias_v)
else:
pos_bias_u = None
pos_bias_v = None
# Positional encodings
self.pos_emb_max_len = pos_emb_max_len
if self_attention_model == "rel_pos":
self.pos_enc = RelPositionalEncoding(
d_model=d_model,
dropout_rate=dropout_pre_encoder,
max_len=pos_emb_max_len,
xscale=self.xscale,
dropout_rate_emb=dropout_emb,
)
elif self_attention_model == "rel_pos_local_attn":
if max(att_context_size) <= 0:
raise ValueError("When using local attention, context size must be set > 0")
self.pos_enc = LocalAttRelPositionalEncoding(
att_context_size=att_context_size,
d_model=d_model,
dropout_rate=dropout,
max_len=pos_emb_max_len,
xscale=self.xscale,
dropout_rate_emb=dropout_emb,
)
elif self_attention_model == "abs_pos":
pos_bias_u = None
pos_bias_v = None
self.pos_enc = PositionalEncoding(
d_model=d_model, dropout_rate=dropout_pre_encoder, max_len=pos_emb_max_len, xscale=self.xscale
)
else:
raise ValueError(f"Not valid self_attention_model: '{self_attention_model}'!")
self.layers = nn.ModuleList()
for i in range(n_layers):
layer = ConformerLayer(
d_model=d_model,
d_ff=d_ff,
self_attention_model=self_attention_model,
global_tokens=global_tokens,
global_tokens_spacing=global_tokens_spacing,
global_attn_separate=global_attn_separate,
n_heads=n_heads,
conv_kernel_size=conv_kernel_size,
conv_norm_type=conv_norm_type,
conv_context_size=self.conv_context_size,
dropout=dropout,
dropout_att=dropout_att,
pos_bias_u=pos_bias_u,
pos_bias_v=pos_bias_v,
att_context_size=self.att_context_size,
use_bias=use_bias,
use_pytorch_sdpa=self.use_pytorch_sdpa,
use_pytorch_sdpa_backends=self.use_pytorch_sdpa_backends,
)
self.layers.append(layer)
if feat_out > 0 and feat_out != self._feat_out:
self.out_proj = nn.Linear(self._feat_out, feat_out)
self._feat_out = feat_out
else:
self.out_proj = None
self._feat_out = d_model
self.set_max_audio_length(self.pos_emb_max_len)
self.use_pad_mask = True
self.setup_streaming_params()
self.export_cache_support = False
self.layer_drop_probs = compute_stochastic_depth_drop_probs(
len(self.layers), stochastic_depth_drop_prob, stochastic_depth_mode, stochastic_depth_start_layer
)
# will be set in self.forward() if defined in AccessMixin config
self.interctc_capture_at_layers = None
@typecheck()
def forward(
self,
audio_signal,
length,
cache_last_channel=None,
cache_last_time=None,
cache_last_channel_len=None,
bypass_pre_encode=False,
):
"""
Forward function for the ConformerEncoder accepting an audio signal and its corresponding length.
The `audio_signal` input supports two formats depending on the `bypass_pre_encode` boolean flag.
This determines the required format of the input variable `audio_signal`:
(1) bypass_pre_encode = False (default):
`audio_signal` must be a tensor containing audio features.
Shape: (batch, self._feat_in, n_frames)
(2) bypass_pre_encode = True:
`audio_signal` must be a tensor containing pre-encoded embeddings.
Shape: (batch, n_frame, self.d_model)
"""
if not bypass_pre_encode and audio_signal.shape[-2] != self._feat_in:
raise ValueError(
f"If bypass_pre_encode is False, audio_signal should have shape "
f"(batch, {self._feat_in}, n_frame) but got last dimension {audio_signal.shape[-2]}."
)
if bypass_pre_encode and audio_signal.shape[-1] != self.d_model:
raise ValueError(
f"If bypass_pre_encode is True, audio_signal should have shape "
f"(batch, n_frame, {self.d_model}) but got last dimension {audio_signal.shape[-1]}."
)
if bypass_pre_encode:
self.update_max_seq_length(
seq_length=audio_signal.size(2) * self.subsampling_factor, device=audio_signal.device
)
else:
self.update_max_seq_length(seq_length=audio_signal.size(2), device=audio_signal.device)
return self.forward_internal(
audio_signal,
length,
cache_last_channel=cache_last_channel,
cache_last_time=cache_last_time,
cache_last_channel_len=cache_last_channel_len,
bypass_pre_encode=bypass_pre_encode,
)
def forward_internal(
self,
audio_signal,
length,
cache_last_channel=None,
cache_last_time=None,
cache_last_channel_len=None,
bypass_pre_encode=False,
):
"""
The `audio_signal` input supports two formats depending on the `bypass_pre_encode` boolean flag.
This determines the required format of the input variable `audio_signal`:
(1) bypass_pre_encode = False (default):
`audio_signal` must be a tensor containing audio features.
Shape: (batch, self._feat_in, n_frames)
(2) bypass_pre_encode = True:
`audio_signal` must be a tensor containing pre-encoded embeddings.
Shape: (batch, n_frame, self.d_model)
`bypass_pre_encode=True` is used in cases where frame-level, context-independent embeddings are
needed to be saved or reused (e.g., speaker cache in streaming speaker diarization).
"""
if length is None:
length = audio_signal.new_full(
(audio_signal.size(0),), audio_signal.size(-1), dtype=torch.int64, device=audio_signal.device
)
# select a random att_context_size with the distribution specified by att_context_probs during training
# for non-validation cases like test, validation or inference, it uses the first mode in self.att_context_size
if self.training and len(self.att_context_size_all) > 1:
cur_att_context_size = random.choices(self.att_context_size_all, weights=self.att_context_probs)[0]
else:
cur_att_context_size = self.att_context_size
if not bypass_pre_encode:
audio_signal = torch.transpose(audio_signal, 1, 2)
if isinstance(self.pre_encode, nn.Linear):
audio_signal = self.pre_encode(audio_signal)
else:
audio_signal, length = self.pre_encode(x=audio_signal, lengths=length)
length = length.to(torch.int64)
# `self.streaming_cfg` is set by setup_streaming_cfg(), called in the init
if self.streaming_cfg.drop_extra_pre_encoded > 0 and cache_last_channel is not None:
audio_signal = audio_signal[:, self.streaming_cfg.drop_extra_pre_encoded :, :]
length = (length - self.streaming_cfg.drop_extra_pre_encoded).clamp(min=0)
if self.reduction_position is not None and cache_last_channel is not None:
raise ValueError("Caching with reduction feature is not supported yet!")
max_audio_length = audio_signal.size(1)
if cache_last_channel is not None:
cache_len = self.streaming_cfg.last_channel_cache_size
cache_keep_size = max_audio_length - self.streaming_cfg.cache_drop_size
max_audio_length = max_audio_length + cache_len
padding_length = length + cache_len
offset = torch.neg(cache_last_channel_len) + cache_len
else:
padding_length = length
cache_last_channel_next = None
cache_len = 0
offset = None
audio_signal, pos_emb = self.pos_enc(x=audio_signal, cache_len=cache_len)
# Create the self-attention and padding masks
pad_mask, att_mask = self._create_masks(
att_context_size=cur_att_context_size,
padding_length=padding_length,
max_audio_length=max_audio_length,
offset=offset,
device=audio_signal.device,
)
if cache_last_channel is not None:
pad_mask = pad_mask[:, cache_len:]
if att_mask is not None:
att_mask = att_mask[:, cache_len:]
# Convert caches from the tensor to list
cache_last_time_next = []
cache_last_channel_next = []
for lth, (drop_prob, layer) in enumerate(zip(self.layer_drop_probs, self.layers)):
original_signal = audio_signal
if cache_last_channel is not None:
cache_last_channel_cur = cache_last_channel[lth]
cache_last_time_cur = cache_last_time[lth]
else:
cache_last_channel_cur = None
cache_last_time_cur = None
audio_signal = layer(
x=audio_signal,
att_mask=att_mask,
pos_emb=pos_emb,
pad_mask=pad_mask,
cache_last_channel=cache_last_channel_cur,
cache_last_time=cache_last_time_cur,
)
if cache_last_channel_cur is not None:
(audio_signal, cache_last_channel_cur, cache_last_time_cur) = audio_signal
cache_last_channel_next.append(cache_last_channel_cur)
cache_last_time_next.append(cache_last_time_cur)
# applying stochastic depth logic from https://arxiv.org/abs/2102.03216
if self.training and drop_prob > 0.0:
should_drop = torch.rand(1) < drop_prob
# adjusting to match expectation
if should_drop:
# that's not efficient, but it's hard to implement distributed
# version of dropping layers without deadlock or random seed meddling
# so multiplying the signal by 0 to ensure all weights get gradients
audio_signal = audio_signal * 0.0 + original_signal
else:
# not doing this operation if drop prob is 0 as it's identity in that case
audio_signal = (audio_signal - original_signal) / (1.0 - drop_prob) + original_signal
if self.reduction_position == lth:
audio_signal, length = self.reduction_subsampling(x=audio_signal, lengths=length)
max_audio_length = audio_signal.size(1)
# Don't update the audio_signal here because then it will again scale the audio_signal
# and cause an increase in the WER
_, pos_emb = self.pos_enc(x=audio_signal, cache_len=cache_len)
pad_mask, att_mask = self._create_masks(
att_context_size=cur_att_context_size,
padding_length=length,
max_audio_length=max_audio_length,
offset=offset,
device=audio_signal.device,
)
# saving tensors if required for interctc loss
if self.is_access_enabled(getattr(self, "model_guid", None)):
if self.interctc_capture_at_layers is None:
self.interctc_capture_at_layers = self.access_cfg.get("interctc", {}).get("capture_layers", [])
if lth in self.interctc_capture_at_layers:
lth_audio_signal = audio_signal
if self.out_proj is not None:
lth_audio_signal = self.out_proj(audio_signal)
# shape is the same as the shape of audio_signal output, i.e. [B, D, T]
self.register_accessible_tensor(
name=f"interctc/layer_output_{lth}", tensor=torch.transpose(lth_audio_signal, 1, 2)
)
self.register_accessible_tensor(name=f"interctc/layer_length_{lth}", tensor=length)
if self.out_proj is not None:
audio_signal = self.out_proj(audio_signal)
# Reduction
if self.reduction_position == -1:
audio_signal, length = self.reduction_subsampling(x=audio_signal, lengths=length)
audio_signal = torch.transpose(audio_signal, 1, 2)
length = length.to(dtype=torch.int64)
if cache_last_channel is not None:
cache_last_channel_next = torch.stack(cache_last_channel_next, dim=0)
cache_last_time_next = torch.stack(cache_last_time_next, dim=0)
return (
audio_signal,
length,
cache_last_channel_next,
cache_last_time_next,
torch.clamp(cache_last_channel_len + cache_keep_size, max=cache_len),
)
else:
return audio_signal, length
def update_max_seq_length(self, seq_length: int, device):
"""
Updates the maximum sequence length for the model.
Args:
seq_length (int): New maximum sequence length.
device (torch.device): Device to use for computations.
"""
# Find global max audio length across all nodes
if self.sync_max_audio_length and torch.distributed.is_initialized():
global_max_len = torch.tensor([seq_length], dtype=torch.float32, device=device)
# Update across all ranks in the distributed system
torch.distributed.all_reduce(global_max_len, op=torch.distributed.ReduceOp.MAX)
seq_length = global_max_len.int().item()
if seq_length > self.max_audio_length:
self.set_max_audio_length(seq_length)
def set_max_audio_length(self, max_audio_length):
"""
Sets maximum input length.
Pre-calculates internal seq_range mask.
Args:
max_audio_length (int): New maximum sequence length.
"""
self.max_audio_length = max_audio_length
device = next(self.parameters()).device
dtype = next(self.parameters()).dtype
self.pos_enc.extend_pe(max_audio_length, device, dtype)
def _create_masks(self, att_context_size, padding_length, max_audio_length, offset, device):
if self.self_attention_model != "rel_pos_local_attn":
att_mask = torch.ones(1, max_audio_length, max_audio_length, dtype=torch.bool, device=device)
if self.att_context_style == "regular":
if att_context_size[0] >= 0:
att_mask = att_mask.triu(diagonal=-att_context_size[0])
if att_context_size[1] >= 0:
att_mask = att_mask.tril(diagonal=att_context_size[1])
elif self.att_context_style == "chunked_limited":
# When right context is unlimited, just the left side of the masking need to get updated
if att_context_size[1] == -1:
if att_context_size[0] >= 0:
att_mask = att_mask.triu(diagonal=-att_context_size[0])
else:
chunk_size = att_context_size[1] + 1
# left_chunks_num specifies the number of chunks to be visible by each chunk on the left side
if att_context_size[0] >= 0:
left_chunks_num = att_context_size[0] // chunk_size
else:
left_chunks_num = 10000
chunk_idx = torch.arange(0, max_audio_length, dtype=torch.int, device=att_mask.device)
chunk_idx = torch.div(chunk_idx, chunk_size, rounding_mode="trunc")
diff_chunks = chunk_idx.unsqueeze(1) - chunk_idx.unsqueeze(0)
chunked_limited_mask = torch.logical_and(
torch.le(diff_chunks, left_chunks_num), torch.ge(diff_chunks, 0)
)
att_mask = torch.logical_and(att_mask, chunked_limited_mask.unsqueeze(0))
else:
att_mask = None
# pad_mask is the masking to be used to ignore paddings
pad_mask = torch.arange(0, max_audio_length, device=device).expand(
padding_length.size(0), -1
) < padding_length.unsqueeze(-1)
if offset is not None:
pad_mask_off = torch.arange(0, max_audio_length, device=device).expand(
padding_length.size(0), -1
) >= offset.unsqueeze(-1)
pad_mask = pad_mask_off.logical_and(pad_mask)
if att_mask is not None:
# pad_mask_for_att_mask is the mask which helps to ignore paddings
pad_mask_for_att_mask = pad_mask.unsqueeze(1).repeat([1, max_audio_length, 1])
pad_mask_for_att_mask = torch.logical_and(pad_mask_for_att_mask, pad_mask_for_att_mask.transpose(1, 2))
# att_mask is the masking to be used by the MHA layers to ignore the tokens not supposed to be visible
att_mask = att_mask[:, :max_audio_length, :max_audio_length]
# paddings should also get ignored, so pad_mask_for_att_mask is used to ignore their corresponding scores
att_mask = torch.logical_and(pad_mask_for_att_mask, att_mask.to(pad_mask_for_att_mask.device))
att_mask = ~att_mask
pad_mask = ~pad_mask
return pad_mask, att_mask
def _calc_context_sizes(
self, att_context_size, att_context_probs, att_context_style, conv_context_size, conv_kernel_size
):
# convert att_context_size to a standard list of lists
if att_context_size:
att_context_size_all = list(att_context_size)
if isinstance(att_context_size_all[0], int):
att_context_size_all = [att_context_size_all]
for i, att_cs in enumerate(att_context_size_all):
if isinstance(att_cs, ListConfig):
att_context_size_all[i] = list(att_cs)
if att_context_style == "chunked_limited":
if att_cs[0] > 0 and att_cs[0] % (att_cs[1] + 1) > 0:
raise ValueError(f"att_context_size[{i}][0] % (att_context_size[{i}][1] + 1) should be zero!")
if att_cs[1] < 0 and len(att_context_size_all) <= 1:
raise ValueError(
f"Right context (att_context_size[{i}][1]) can not be unlimited for chunked_limited style!"
)
else:
att_context_size_all = [[-1, -1]]
if att_context_probs:
if len(att_context_probs) != len(att_context_size_all):
raise ValueError("The size of the att_context_probs should be the same as att_context_size.")
att_context_probs = list(att_context_probs)
if sum(att_context_probs) != 1:
raise ValueError("The sum of numbers in att_context_probs should be equal to one to be a distribution.")
else:
att_context_probs = [1.0 / len(att_context_size_all)] * len(att_context_size_all)
if conv_context_size is not None:
if isinstance(conv_context_size, ListConfig):
conv_context_size = list(conv_context_size)
if not isinstance(conv_context_size, list) and not isinstance(conv_context_size, str):
raise ValueError(
"Invalid conv_context_size! It should be the string 'causal' or a list of two integers."
)
if conv_context_size == "causal":
conv_context_size = [conv_kernel_size - 1, 0]
else:
if conv_context_size[0] + conv_context_size[1] + 1 != conv_kernel_size:
raise ValueError(f"Invalid conv_context_size: {self.conv_context_size}!")
else:
conv_context_size = [(conv_kernel_size - 1) // 2, (conv_kernel_size - 1) // 2]
return att_context_size_all, att_context_size_all[0], att_context_probs, conv_context_size
def setup_streaming_params(
self,
chunk_size: int = None,
shift_size: int = None,
left_chunks: int = None,
att_context_size: list = None,
max_context: int = 10000,
):
"""
This function sets the needed values and parameters to perform streaming.
The configuration would be stored in self.streaming_cfg.
The streaming configuration is needed to simulate streaming inference.
Args:
chunk_size (int): overrides the chunk size
shift_size (int): overrides the shift size for chunks
left_chunks (int): overrides the number of left chunks visible to each chunk
max_context (int): the value used for the cache size of last_channel layers
if left context is set to infinity (-1)
Defaults to -1 (means feat_out is d_model)
"""
streaming_cfg = CacheAwareStreamingConfig()
# When att_context_size is not specified, it uses the default_att_context_size
if att_context_size is None:
att_context_size = self.att_context_size
if chunk_size is not None:
if chunk_size < 1:
raise ValueError("chunk_size needs to be a number larger or equal to one.")
lookahead_steps = chunk_size - 1
streaming_cfg.cache_drop_size = chunk_size - shift_size
elif self.att_context_style == "chunked_limited":
lookahead_steps = att_context_size[1]
streaming_cfg.cache_drop_size = 0
elif self.att_context_style == "regular":
lookahead_steps = att_context_size[1] * self.n_layers + self.conv_context_size[1] * self.n_layers
streaming_cfg.cache_drop_size = lookahead_steps
else:
streaming_cfg.cache_drop_size = 0
lookahead_steps = None
if chunk_size is None:
streaming_cfg.last_channel_cache_size = att_context_size[0] if att_context_size[0] >= 0 else max_context
else:
if left_chunks is None:
streaming_cfg.last_channel_cache_size = att_context_size[0] if att_context_size[0] >= 0 else max_context
logging.warning(
f"left_chunks is not set. Setting it to default: {streaming_cfg.last_channel_cache_size}."
)
else:
streaming_cfg.last_channel_cache_size = left_chunks * chunk_size
if hasattr(self.pre_encode, "get_sampling_frames"):
sampling_frames = self.pre_encode.get_sampling_frames()
else:
sampling_frames = 0
if isinstance(sampling_frames, list):
streaming_cfg.chunk_size = [
sampling_frames[0] + self.subsampling_factor * lookahead_steps,
sampling_frames[1] + self.subsampling_factor * lookahead_steps,
]
else:
streaming_cfg.chunk_size = sampling_frames * (1 + lookahead_steps)
if isinstance(sampling_frames, list):
streaming_cfg.shift_size = [
sampling_frames[0] + sampling_frames[1] * (lookahead_steps - streaming_cfg.cache_drop_size),
sampling_frames[1] + sampling_frames[1] * (lookahead_steps - streaming_cfg.cache_drop_size),
]
else:
streaming_cfg.shift_size = sampling_frames * (1 + lookahead_steps - streaming_cfg.cache_drop_size)
if isinstance(streaming_cfg.shift_size, list):
streaming_cfg.valid_out_len = (
streaming_cfg.shift_size[1] - sampling_frames[1]
) // self.subsampling_factor + 1
else:
streaming_cfg.valid_out_len = streaming_cfg.shift_size // self.subsampling_factor
if hasattr(self.pre_encode, "get_streaming_cache_size"):
streaming_cfg.pre_encode_cache_size = self.pre_encode.get_streaming_cache_size()
else:
streaming_cfg.pre_encode_cache_size = 0
if isinstance(streaming_cfg.pre_encode_cache_size, list):
if streaming_cfg.pre_encode_cache_size[1] >= 1:
streaming_cfg.drop_extra_pre_encoded = (
1 + (streaming_cfg.pre_encode_cache_size[1] - 1) // self.subsampling_factor
)
else:
streaming_cfg.drop_extra_pre_encoded = 0
else:
streaming_cfg.drop_extra_pre_encoded = streaming_cfg.pre_encode_cache_size // self.subsampling_factor
for m in self.layers.modules():
if hasattr(m, "_max_cache_len"):
if isinstance(m, MultiHeadAttention):
m.cache_drop_size = streaming_cfg.cache_drop_size
if isinstance(m, CausalConv1D):
m.cache_drop_size = streaming_cfg.cache_drop_size
self.streaming_cfg = streaming_cfg
def change_subsampling_conv_chunking_factor(self, subsampling_conv_chunking_factor: int):
"""
Update the conv_chunking_factor (int)
Default is 1 (auto)
Set it to -1 (disabled) or to a specific value (power of 2) if you OOM in the conv subsampling layers
Args:
subsampling_conv_chunking_factor (int)
"""
if not hasattr(self.pre_encode, "change_subsampling_conv_chunking_factor"):
logging.info("Model pre_encoder doesn't have a change_subsampling_conv_chunking_factor method ")
return
self.pre_encode.change_subsampling_conv_chunking_factor(
subsampling_conv_chunking_factor=subsampling_conv_chunking_factor
)
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_modules.py
+372
View File
@@ -0,0 +1,372 @@
# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/submodules/conformer_modules.py; training-only code removed.
import torch
from torch import nn as nn
from torch.nn import LayerNorm
from ..compat import AccessMixin, AttentionAdapterModuleMixin, FusedBatchNorm1d, Swish, activation_registry
from .causal_convs import CausalConv1D
from .multi_head_attention import (
MultiHeadAttention,
RelPositionMultiHeadAttention,
RelPositionMultiHeadAttentionLongformer,
)
__all__ = ["ConformerConvolution", "ConformerFeedForward", "ConformerLayer"]
class ConformerLayer(torch.nn.Module, AttentionAdapterModuleMixin, AccessMixin):
"""A single block of the Conformer encoder.
Args:
d_model (int): input dimension of MultiheadAttentionMechanism and PositionwiseFeedForward
d_ff (int): hidden dimension of PositionwiseFeedForward
self_attention_model (str): type of the attention layer and positional encoding
'rel_pos': relative positional embedding and Transformer-XL
'rel_pos_local_attn': relative positional embedding and Transformer-XL with local attention using
overlapping chunks. Attention context is determined by att_context_size parameter.
'abs_pos': absolute positional embedding and Transformer
Default is rel_pos.
global_tokens (int): number of tokens to be used for global attention.
Only relevant if self_attention_model is 'rel_pos_local_attn'.
Defaults to 0.
global_tokens_spacing (int): how far apart the global tokens are
Defaults to 1.
global_attn_separate (bool): whether the q, k, v layers used for global tokens should be separate.
Defaults to False.
n_heads (int): number of heads for multi-head attention
conv_kernel_size (int): kernel size for depthwise convolution in convolution module
dropout (float): dropout probabilities for linear layers
dropout_att (float): dropout probabilities for attention distributions
use_bias (bool): Apply bias to all Linear and Conv1d layers from each ConformerLayer to improve activation flow and stabilize training of huge models.
Defaults to True.
"""
def __init__(
self,
d_model,
d_ff,
self_attention_model="rel_pos",
global_tokens=0,
global_tokens_spacing=1,
global_attn_separate=False,
n_heads=4,
conv_kernel_size=31,
conv_norm_type="batch_norm",
conv_context_size=None,
dropout=0.1,
dropout_att=0.1,
pos_bias_u=None,
pos_bias_v=None,
att_context_size=[-1, -1],
use_bias=True,
use_pytorch_sdpa=False,
use_pytorch_sdpa_backends=None,
):
super(ConformerLayer, self).__init__()
self.use_pytorch_sdpa = use_pytorch_sdpa
if use_pytorch_sdpa_backends is None:
use_pytorch_sdpa_backends = []
self.use_pytorch_sdpa_backends = use_pytorch_sdpa_backends
self.self_attention_model = self_attention_model
self.n_heads = n_heads
self.fc_factor = 0.5
# first feed forward module
self.norm_feed_forward1 = LayerNorm(d_model)
self.feed_forward1 = ConformerFeedForward(d_model=d_model, d_ff=d_ff, dropout=dropout, use_bias=use_bias)
# convolution module
self.norm_conv = LayerNorm(d_model)
self.conv = ConformerConvolution(
d_model=d_model,
kernel_size=conv_kernel_size,
norm_type=conv_norm_type,
conv_context_size=conv_context_size,
use_bias=use_bias,
)
# multi-headed self-attention module
self.norm_self_att = LayerNorm(d_model)
MHA_max_cache_len = att_context_size[0]
if self_attention_model == "rel_pos":
self.self_attn = RelPositionMultiHeadAttention(
n_head=n_heads,
n_feat=d_model,
dropout_rate=dropout_att,
pos_bias_u=pos_bias_u,
pos_bias_v=pos_bias_v,
max_cache_len=MHA_max_cache_len,
use_bias=use_bias,
use_pytorch_sdpa=self.use_pytorch_sdpa,
use_pytorch_sdpa_backends=self.use_pytorch_sdpa_backends,
)
elif self_attention_model == "rel_pos_local_attn":
self.self_attn = RelPositionMultiHeadAttentionLongformer(
n_head=n_heads,
n_feat=d_model,
dropout_rate=dropout_att,
pos_bias_u=pos_bias_u,
pos_bias_v=pos_bias_v,
max_cache_len=MHA_max_cache_len,
att_context_size=att_context_size,
global_tokens=global_tokens,
global_tokens_spacing=global_tokens_spacing,
global_attn_separate=global_attn_separate,
use_bias=use_bias,
)
elif self_attention_model == "abs_pos":
self.self_attn = MultiHeadAttention(
n_head=n_heads,
n_feat=d_model,
dropout_rate=dropout_att,
max_cache_len=MHA_max_cache_len,
use_bias=use_bias,
use_pytorch_sdpa=self.use_pytorch_sdpa,
use_pytorch_sdpa_backends=self.use_pytorch_sdpa_backends,
)
else:
raise ValueError(
f"'{self_attention_model}' is not not a valid value for 'self_attention_model', "
f"valid values can be from ['rel_pos', 'rel_pos_local_attn', 'abs_pos']"
)
# second feed forward module
self.norm_feed_forward2 = LayerNorm(d_model)
self.feed_forward2 = ConformerFeedForward(d_model=d_model, d_ff=d_ff, dropout=dropout, use_bias=use_bias)
self.dropout = nn.Dropout(dropout)
self.norm_out = LayerNorm(d_model)
def forward(self, x, att_mask=None, pos_emb=None, pad_mask=None, cache_last_channel=None, cache_last_time=None):
"""
Args:
x (torch.Tensor): input signals (B, T, d_model)
att_mask (torch.Tensor): attention masks(B, T, T)
pos_emb (torch.Tensor): (L, 1, d_model)
pad_mask (torch.tensor): padding mask
cache_last_channel (torch.tensor) : cache for MHA layers (B, T_cache, d_model)
cache_last_time (torch.tensor) : cache for convolutional layers (B, d_model, T_cache)
Returns:
x (torch.Tensor): (B, T, d_model)
cache_last_channel (torch.tensor) : next cache for MHA layers (B, T_cache, d_model)
cache_last_time (torch.tensor) : next cache for convolutional layers (B, d_model, T_cache)
"""
residual = x
x = self.norm_feed_forward1(x)
x = self.feed_forward1(x)
residual = residual + self.dropout(x) * self.fc_factor
x = self.norm_self_att(residual)
if self.self_attention_model == "rel_pos":
x = self.self_attn(query=x, key=x, value=x, mask=att_mask, pos_emb=pos_emb, cache=cache_last_channel)
elif self.self_attention_model == "rel_pos_local_attn":
x = self.self_attn(query=x, key=x, value=x, pad_mask=pad_mask, pos_emb=pos_emb, cache=cache_last_channel)
elif self.self_attention_model == "abs_pos":
x = self.self_attn(query=x, key=x, value=x, mask=att_mask, cache=cache_last_channel)
else:
x = None
if x is not None and cache_last_channel is not None:
(x, cache_last_channel) = x
residual = residual + self.dropout(x)
if self.is_adapter_available():
# Call the MHA adapters
pack_input = {
"x": residual,
"loc": "mha",
"att_mask": att_mask,
"pos_emb": pos_emb,
}
pack_input = self.forward_enabled_adapters(pack_input)
residual = pack_input["x"]
x = self.norm_conv(residual)
x = self.conv(x, pad_mask=pad_mask, cache=cache_last_time)
if cache_last_time is not None:
(x, cache_last_time) = x
residual = residual + self.dropout(x)
x = self.norm_feed_forward2(residual)
x = self.feed_forward2(x)
residual = residual + self.dropout(x) * self.fc_factor
x = self.norm_out(residual)
if self.is_adapter_available():
# Call the adapters
pack_input = {
"x": x,
"loc": "post",
}
pack_input = self.forward_enabled_adapters(pack_input)
x = pack_input["x"]
if self.is_access_enabled(getattr(self, "model_guid", None)) and self.access_cfg.get(
"save_encoder_tensors", False
):
self.register_accessible_tensor(name="encoder", tensor=x)
if cache_last_channel is None:
return x
else:
return x, cache_last_channel, cache_last_time
class ConformerConvolution(nn.Module):
"""The convolution module for the Conformer model.
Args:
d_model (int): hidden dimension
kernel_size (int): kernel size for depthwise convolution
pointwise_activation (str): name of the activation function to be used for the pointwise conv.
Note that Conformer uses a special key `glu_` which is treated as the original default from
the paper.
use_bias (bool): Use bias in all Linear and Conv1d layers improve activation flow and stabilize training of huge models.
Defaults to True
"""
def __init__(
self,
d_model,
kernel_size,
norm_type="batch_norm",
conv_context_size=None,
pointwise_activation="glu_",
use_bias=True,
):
super(ConformerConvolution, self).__init__()
assert (kernel_size - 1) % 2 == 0
self.d_model = d_model
self.kernel_size = kernel_size
self.norm_type = norm_type
self.use_bias = use_bias
if conv_context_size is None:
conv_context_size = (kernel_size - 1) // 2
if pointwise_activation in activation_registry:
self.pointwise_activation = activation_registry[pointwise_activation]()
dw_conv_input_dim = d_model * 2
if hasattr(self.pointwise_activation, "inplace"):
self.pointwise_activation.inplace = True
else:
self.pointwise_activation = pointwise_activation
dw_conv_input_dim = d_model
self.pointwise_conv1 = nn.Conv1d(
in_channels=d_model,
out_channels=d_model * 2,
kernel_size=1,
stride=1,
padding=0,
bias=self.use_bias,
)
self.depthwise_conv = CausalConv1D(
in_channels=dw_conv_input_dim,
out_channels=dw_conv_input_dim,
kernel_size=kernel_size,
stride=1,
padding=conv_context_size,
groups=dw_conv_input_dim,
bias=self.use_bias,
)
if norm_type == "batch_norm":
self.batch_norm = nn.BatchNorm1d(dw_conv_input_dim)
elif norm_type == "instance_norm":
self.batch_norm = nn.InstanceNorm1d(dw_conv_input_dim)
elif norm_type == "layer_norm":
self.batch_norm = nn.LayerNorm(dw_conv_input_dim)
elif norm_type == "fused_batch_norm":
self.batch_norm = FusedBatchNorm1d(dw_conv_input_dim)
elif norm_type.startswith("group_norm"):
num_groups = int(norm_type.replace("group_norm", ""))
self.batch_norm = nn.GroupNorm(num_groups=num_groups, num_channels=d_model)
else:
raise ValueError(f"conv_norm_type={norm_type} is not valid!")
self.activation = Swish()
self.pointwise_conv2 = nn.Conv1d(
in_channels=dw_conv_input_dim,
out_channels=d_model,
kernel_size=1,
stride=1,
padding=0,
bias=self.use_bias,
)
def forward(self, x, pad_mask=None, cache=None):
x = x.transpose(1, 2)
x = self.pointwise_conv1(x)
# Compute the activation function or use GLU for original Conformer
if self.pointwise_activation == "glu_":
x = nn.functional.glu(x, dim=1)
else:
x = self.pointwise_activation(x)
if pad_mask is not None:
x = x.masked_fill(pad_mask.unsqueeze(1), 0.0)
x = self.depthwise_conv(x, cache=cache)
if cache is not None:
x, cache = x
if self.norm_type == "layer_norm":
x = x.transpose(1, 2)
x = self.batch_norm(x)
x = x.transpose(1, 2)
else:
x = self.batch_norm(x)
x = self.activation(x)
x = self.pointwise_conv2(x)
x = x.transpose(1, 2)
if cache is None:
return x
else:
return x, cache
class ConformerFeedForward(nn.Module):
"""
feed-forward module of Conformer model.
use_bias (bool): Apply bias to all Linear and Conv1d layers improve activation flow and stabilize training of huge models.
"""
def __init__(self, d_model, d_ff, dropout, activation=Swish(), use_bias=True):
super(ConformerFeedForward, self).__init__()
self.d_model = d_model
self.d_ff = d_ff
self.use_bias = use_bias
self.linear1 = nn.Linear(d_model, d_ff, bias=self.use_bias)
self.activation = activation
self.dropout = nn.Dropout(p=dropout)
self.linear2 = nn.Linear(d_ff, d_model, bias=self.use_bias)
def forward(self, x):
x = self.linear1(x)
x = self.activation(x)
x = self.dropout(x)
x = self.linear2(x)
return x
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/features.py
+633
View File
@@ -0,0 +1,633 @@
# Copyright (c) 2018 Ryan Leary
# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# This file contains code artifacts adapted from https://github.com/ryanleary/patter
# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/preprocessing/features.py; training-only code removed.
# Removed: WaveformFeaturizer / FeaturizerFactory (data-loading only, required perturb/segment).
import math
import random
from typing import Optional, Tuple, Union
import numpy as np
import torch
import torch.nn as nn
from ..compat import logging
try:
import torchaudio
HAVE_TORCHAUDIO = True
except ModuleNotFoundError:
HAVE_TORCHAUDIO = False
CONSTANT = 1e-5
def normalize_batch(x, seq_len, normalize_type):
x_mean = None
x_std = None
if normalize_type == "per_feature":
batch_size = x.shape[0]
max_time = x.shape[2]
# When doing stream capture to a graph, item() is not allowed
# because it calls cudaStreamSynchronize(). Therefore, we are
# sacrificing some error checking when running with cuda graphs.
if (
torch.cuda.is_available()
and not torch.cuda.is_current_stream_capturing()
and torch.any(seq_len == 1).item()
):
raise ValueError(
"normalize_batch with `per_feature` normalize_type received a tensor of length 1. This will result "
"in torch.std() returning nan. Make sure your audio length has enough samples for a single "
"feature (ex. at least `hop_length` for Mel Spectrograms)."
)
time_steps = torch.arange(max_time, device=x.device).unsqueeze(0).expand(batch_size, max_time)
valid_mask = time_steps < seq_len.unsqueeze(1)
x_mean_numerator = torch.where(valid_mask.unsqueeze(1), x, 0.0).sum(axis=2)
x_mean_denominator = valid_mask.sum(axis=1)
x_mean = x_mean_numerator / x_mean_denominator.unsqueeze(1)
# Subtract 1 in the denominator to correct for the bias.
x_std = torch.sqrt(
torch.sum(torch.where(valid_mask.unsqueeze(1), x - x_mean.unsqueeze(2), 0.0) ** 2, axis=2)
/ (x_mean_denominator.unsqueeze(1) - 1.0)
)
# make sure x_std is not zero
x_std += CONSTANT
return (x - x_mean.unsqueeze(2)) / x_std.unsqueeze(2), x_mean, x_std
elif normalize_type == "all_features":
x_mean = torch.zeros(seq_len.shape, dtype=x.dtype, device=x.device)
x_std = torch.zeros(seq_len.shape, dtype=x.dtype, device=x.device)
for i in range(x.shape[0]):
x_mean[i] = x[i, :, : seq_len[i].item()].mean()
x_std[i] = x[i, :, : seq_len[i].item()].std()
# make sure x_std is not zero
x_std += CONSTANT
return (x - x_mean.view(-1, 1, 1)) / x_std.view(-1, 1, 1), x_mean, x_std
elif "fixed_mean" in normalize_type and "fixed_std" in normalize_type:
x_mean = torch.tensor(normalize_type["fixed_mean"], device=x.device)
x_std = torch.tensor(normalize_type["fixed_std"], device=x.device)
return (
(x - x_mean.view(x.shape[0], x.shape[1]).unsqueeze(2)) / x_std.view(x.shape[0], x.shape[1]).unsqueeze(2),
x_mean,
x_std,
)
else:
return x, x_mean, x_std
def splice_frames(x, frame_splicing):
"""Stacks frames together across feature dim
input is batch_size, feature_dim, num_frames
output is batch_size, feature_dim*frame_splicing, num_frames
"""
seq = [x]
for n in range(1, frame_splicing):
seq.append(torch.cat([x[:, :, :n], x[:, :, n:]], dim=2))
return torch.cat(seq, dim=1)
@torch.jit.script_if_tracing
def make_seq_mask_like(
lengths: torch.Tensor, like: torch.Tensor, time_dim: int = -1, valid_ones: bool = True
) -> torch.Tensor:
"""
Args:
lengths: Tensor with shape [B] containing the sequence length of each batch element
like: The mask will contain the same number of dimensions as this Tensor, and will have the same max
length in the time dimension of this Tensor.
time_dim: Time dimension of the `shape_tensor` and the resulting mask. Zero-based.
valid_ones: If True, valid tokens will contain value `1` and padding will be `0`. Else, invert.
Returns:
A :class:`torch.Tensor` containing 1's and 0's for valid and invalid tokens, respectively, if `valid_ones`, else
vice-versa. Mask will have the same number of dimensions as `like`. Batch and time dimensions will match
the `like`. All other dimensions will be singletons. E.g., if `like.shape == [3, 4, 5]` and
`time_dim == -1', mask will have shape `[3, 1, 5]`.
"""
# Mask with shape [B, T]
mask = torch.arange(like.shape[time_dim], device=like.device).repeat(lengths.shape[0], 1).lt(lengths.view(-1, 1))
# [B, T] -> [B, *, T] where * is any number of singleton dimensions to expand to like tensor
for _ in range(like.dim() - mask.dim()):
mask = mask.unsqueeze(1)
# If needed, transpose time dim
if time_dim != -1 and time_dim != mask.dim() - 1:
mask = mask.transpose(-1, time_dim)
# Maybe invert the padded vs. valid token values
if not valid_ones:
mask = ~mask
return mask
class FilterbankFeatures(nn.Module):
"""Featurizer that converts wavs to Mel Spectrograms.
See AudioToMelSpectrogramPreprocessor for args.
"""
def __init__(
self,
sample_rate=16000,
n_window_size=320,
n_window_stride=160,
window="hann",
normalize="per_feature",
n_fft=None,
preemph=0.97,
nfilt=64,
lowfreq=0,
highfreq=None,
log=True,
log_zero_guard_type="add",
log_zero_guard_value=2**-24,
dither=CONSTANT,
pad_to=16,
max_duration=16.7,
frame_splicing=1,
exact_pad=False,
pad_value=0,
mag_power=2.0,
use_grads=False,
rng=None,
nb_augmentation_prob=0.0,
nb_max_freq=4000,
mel_norm="slaney",
stft_exact_pad=False, # Deprecated arguments; kept for config compatibility
stft_conv=False, # Deprecated arguments; kept for config compatibility
):
super().__init__()
if stft_conv or stft_exact_pad:
logging.warning(
"Using torch_stft is deprecated and has been removed. The values have been forcibly set to False "
"for FilterbankFeatures and AudioToMelSpectrogramPreprocessor. Please set exact_pad to True "
"as needed."
)
if exact_pad and n_window_stride % 2 == 1:
raise NotImplementedError(
f"{self} received exact_pad == True, but hop_size was odd. If audio_length % hop_size == 0. Then the "
"returned spectrogram would not be of length audio_length // hop_size. Please use an even hop_size."
)
self.log_zero_guard_value = log_zero_guard_value
if (
n_window_size is None
or n_window_stride is None
or not isinstance(n_window_size, int)
or not isinstance(n_window_stride, int)
or n_window_size <= 0
or n_window_stride <= 0
):
raise ValueError(
f"{self} got an invalid value for either n_window_size or n_window_stride. Both must be positive ints."
)
logging.info(f"PADDING: {pad_to}")
self.win_length = n_window_size
self.hop_length = n_window_stride
self.n_fft = n_fft or 2 ** math.ceil(math.log2(self.win_length))
self.stft_pad_amount = (self.n_fft - self.hop_length) // 2 if exact_pad else None
self.exact_pad = exact_pad
if exact_pad:
logging.info("STFT using exact pad")
torch_windows = {
"hann": torch.hann_window,
"hamming": torch.hamming_window,
"blackman": torch.blackman_window,
"bartlett": torch.bartlett_window,
"none": None,
}
window_fn = torch_windows.get(window, None)
window_tensor = window_fn(self.win_length, periodic=False) if window_fn else None
self.register_buffer("window", window_tensor)
self.normalize = normalize
self.log = log
self.dither = dither
self.frame_splicing = frame_splicing
self.nfilt = nfilt
self.preemph = preemph
self.pad_to = pad_to
highfreq = highfreq or sample_rate / 2
# librosa.filters.mel replaced by the torchaudio equivalent (norm="slaney" +
# mel_scale="slaney" matches librosa's default HTK=False convention; filterbank
# max abs diff 2.6e-7, verified token- and WAV-identical on the parity fixture).
if not HAVE_TORCHAUDIO:
raise ValueError(f"Need to install torchaudio to instantiate a {self.__class__.__name__}")
filterbanks = (
torchaudio.functional.melscale_fbanks(
n_freqs=self.n_fft // 2 + 1,
f_min=lowfreq,
f_max=highfreq,
n_mels=nfilt,
sample_rate=sample_rate,
norm=mel_norm,
mel_scale="slaney",
)
.T.to(torch.float)
.unsqueeze(0)
)
self.register_buffer("fb", filterbanks)
# Calculate maximum sequence length
max_length = self.get_seq_len(torch.tensor(max_duration * sample_rate, dtype=torch.float))
max_pad = pad_to - (max_length % pad_to) if pad_to > 0 else 0
self.max_length = max_length + max_pad
self.pad_value = pad_value
self.mag_power = mag_power
# We want to avoid taking the log of zero
# There are two options: either adding or clamping to a small value
if log_zero_guard_type not in ["add", "clamp"]:
raise ValueError(
f"{self} received {log_zero_guard_type} for the "
f"log_zero_guard_type parameter. It must be either 'add' or "
f"'clamp'."
)
self.use_grads = use_grads
if not use_grads:
self.forward = torch.no_grad()(self.forward)
self._rng = random.Random() if rng is None else rng
self.nb_augmentation_prob = nb_augmentation_prob
if self.nb_augmentation_prob > 0.0:
if nb_max_freq >= sample_rate / 2:
self.nb_augmentation_prob = 0.0
else:
self._nb_max_fft_bin = int((nb_max_freq / sample_rate) * n_fft)
# log_zero_guard_value is the the small we want to use, we support
# an actual number, or "tiny", or "eps"
self.log_zero_guard_type = log_zero_guard_type
logging.debug(f"sr: {sample_rate}")
logging.debug(f"n_fft: {self.n_fft}")
logging.debug(f"win_length: {self.win_length}")
logging.debug(f"hop_length: {self.hop_length}")
logging.debug(f"n_mels: {nfilt}")
logging.debug(f"fmin: {lowfreq}")
logging.debug(f"fmax: {highfreq}")
logging.debug(f"using grads: {use_grads}")
logging.debug(f"nb_augmentation_prob: {nb_augmentation_prob}")
def stft(self, x):
return torch.stft(
x,
n_fft=self.n_fft,
hop_length=self.hop_length,
win_length=self.win_length,
center=False if self.exact_pad else True,
window=self.window.to(dtype=torch.float, device=x.device),
return_complex=True,
)
def log_zero_guard_value_fn(self, x):
if isinstance(self.log_zero_guard_value, str):
if self.log_zero_guard_value == "tiny":
return torch.finfo(x.dtype).tiny
elif self.log_zero_guard_value == "eps":
return torch.finfo(x.dtype).eps
else:
raise ValueError(
f"{self} received {self.log_zero_guard_value} for the "
f"log_zero_guard_type parameter. It must be either a "
f"number, 'tiny', or 'eps'"
)
else:
return self.log_zero_guard_value
def get_seq_len(self, seq_len):
# Assuming that center is True is stft_pad_amount = 0
pad_amount = self.stft_pad_amount * 2 if self.stft_pad_amount is not None else self.n_fft // 2 * 2
seq_len = torch.floor_divide((seq_len + pad_amount - self.n_fft), self.hop_length) + 1
return seq_len.to(dtype=torch.long)
def forward(self, x, seq_len, linear_spec=False):
# Initialize debug dictionary
# debug_snapshots = {}
# debug_snapshots['step0_input'] = x[0, :1280].cpu() if x.shape[1] >= 1280 else x[0].cpu()
# debug_snapshots['input_len'] = seq_len[0].item()
seq_len = self.get_seq_len(seq_len)
# debug_snapshots['output_seq_len'] = seq_len[0].item()
if self.stft_pad_amount is not None:
x = torch.nn.functional.pad(
x.unsqueeze(1), (self.stft_pad_amount, self.stft_pad_amount), "reflect"
).squeeze(1)
# debug_snapshots['step1_after_reflect_pad'] = x[0, :1280].cpu() if x.shape[1] >= 1280 else x[0].cpu()
# dither (only in training mode for eval determinism)
if self.training and self.dither > 0:
x += self.dither * torch.randn_like(x)
# debug_snapshots['step2_after_dither'] = x[0, :1280].cpu() if x.shape[1] >= 1280 else x[0].cpu()
# do preemphasis
if self.preemph is not None:
x = torch.cat((x[:, 0].unsqueeze(1), x[:, 1:] - self.preemph * x[:, :-1]), dim=1)
# debug_snapshots['step3_after_preemph'] = x[0, :1280].cpu() if x.shape[1] >= 1280 else x[0].cpu()
# disable autocast to get full range of stft values
with torch.amp.autocast(x.device.type, enabled=False):
x = self.stft(x)
# torch stft returns complex tensor (of shape [B,N,T]); so convert to magnitude
# guard is needed for sqrt if grads are passed through
guard = 0 if not self.use_grads else CONSTANT
x = torch.view_as_real(x)
x = torch.sqrt(x.pow(2).sum(-1) + guard)
# debug_snapshots['step4_after_stft_magnitude'] = x[0, :, :10].cpu() # First 10 frames, all freq bins
if self.training and self.nb_augmentation_prob > 0.0:
for idx in range(x.shape[0]):
if self._rng.random() < self.nb_augmentation_prob:
x[idx, self._nb_max_fft_bin :, :] = 0.0
# get power spectrum
if self.mag_power != 1.0:
x = x.pow(self.mag_power)
# debug_snapshots['step5_after_power'] = x[0, :, :10].cpu()
# return plain spectrogram if required
if linear_spec:
return x, seq_len
# disable autocast, otherwise it might be automatically casted to fp16
# on fp16 compatible GPUs and get NaN values for input value of 65520
with torch.amp.autocast(x.device.type, enabled=False):
# dot with filterbank energies
x = torch.matmul(self.fb.to(x.dtype), x)
# debug_snapshots['step6_after_mel_filterbank'] = x[0, :, :10].cpu()
# log features if required
if self.log:
if self.log_zero_guard_type == "add":
x = torch.log(x + self.log_zero_guard_value_fn(x))
elif self.log_zero_guard_type == "clamp":
x = torch.log(torch.clamp(x, min=self.log_zero_guard_value_fn(x)))
else:
raise ValueError("log_zero_guard_type was not understood")
# debug_snapshots['step7_after_log'] = x[0, :, :10].cpu()
# frame splicing if required
if self.frame_splicing > 1:
x = splice_frames(x, self.frame_splicing)
# debug_snapshots['step8_after_frame_splicing'] = x[0, :, :10].cpu()
# normalize if required
if self.normalize:
x, _, _ = normalize_batch(x, seq_len, normalize_type=self.normalize)
# debug_snapshots['step9_after_normalize'] = x[0, :, :10].cpu()
# mask to zero any values beyond seq_len in batch, pad to multiple of `pad_to` (for efficiency)
max_len = x.size(-1)
mask = torch.arange(max_len, device=x.device)
mask = mask.repeat(x.size(0), 1) >= seq_len.unsqueeze(1)
x = x.masked_fill(mask.unsqueeze(1).type(torch.bool).to(device=x.device), self.pad_value)
# debug_snapshots['step10_after_masking'] = x[0, :, :10].cpu()
del mask
pad_to = self.pad_to
if pad_to == "max":
x = nn.functional.pad(x, (0, self.max_length - x.size(-1)), value=self.pad_value)
elif pad_to > 0:
pad_amt = x.size(-1) % pad_to
if pad_amt != 0:
x = nn.functional.pad(x, (0, pad_to - pad_amt), value=self.pad_value)
# debug_snapshots['step11_after_padding'] = x[0, :, :10].cpu()
# Debug: Save all intermediate steps
# import os
# # Add final output and config to debug snapshots
# debug_snapshots['final_output'] = x[0].cpu()
# debug_snapshots['final_shape'] = list(x.shape)
# debug_snapshots['config'] = {
# 'pad_to': pad_to,
# 'normalize': self.normalize,
# 'exact_pad': self.exact_pad,
# 'n_fft': self.n_fft,
# 'hop_length': self.hop_length,
# 'preemph': self.preemph,
# 'mag_power': self.mag_power,
# }
# torch.save(debug_snapshots, debug_path)
# print(f"[Preprocessor Debug] Saved {len(debug_snapshots)} steps to {debug_path}")
# print(f" Input length: {debug_snapshots['input_len']}, Output seq_len: {seq_len[0].item()}, Final shape: {x.shape}")
return x, seq_len
class FilterbankFeaturesTA(nn.Module):
"""
Exportable, `torchaudio`-based implementation of Mel Spectrogram extraction.
See `AudioToMelSpectrogramPreprocessor` for args.
"""
def __init__(
self,
sample_rate: int = 16000,
n_window_size: int = 320,
n_window_stride: int = 160,
normalize: str | None = "per_feature",
nfilt: int = 64,
n_fft: int | None = None,
preemph: float = 0.97,
lowfreq: float = 0,
highfreq: float | None = None,
log: bool = True,
log_zero_guard_type: str = "add",
log_zero_guard_value: Union[float, str] = 2**-24,
dither: float = 1e-5,
window: str = "hann",
pad_to: int = 0,
pad_value: float = 0.0,
mel_norm="slaney",
# Seems like no one uses these options anymore. Don't convolute the code by supporting them.
use_grads: bool = False, # Deprecated arguments; kept for config compatibility
max_duration: float = 16.7, # Deprecated arguments; kept for config compatibility
frame_splicing: int = 1, # Deprecated arguments; kept for config compatibility
exact_pad: bool = False, # Deprecated arguments; kept for config compatibility
nb_augmentation_prob: float = 0.0, # Deprecated arguments; kept for config compatibility
nb_max_freq: int = 4000, # Deprecated arguments; kept for config compatibility
mag_power: float = 2.0, # Deprecated arguments; kept for config compatibility
rng: random.Random | None = None, # Deprecated arguments; kept for config compatibility
stft_exact_pad: bool = False, # Deprecated arguments; kept for config compatibility
stft_conv: bool = False, # Deprecated arguments; kept for config compatibility
):
super().__init__()
if not HAVE_TORCHAUDIO:
raise ValueError(f"Need to install torchaudio to instantiate a {self.__class__.__name__}")
# Make sure log zero guard is supported, if given as a string
supported_log_zero_guard_strings = {"eps", "tiny"}
if isinstance(log_zero_guard_value, str) and log_zero_guard_value not in supported_log_zero_guard_strings:
raise ValueError(
f"Log zero guard value must either be a float or a member of {supported_log_zero_guard_strings}"
)
# Copied from `AudioPreprocessor` due to the ad-hoc structuring of the Mel Spec extractor class
self.torch_windows = {
"hann": torch.hann_window,
"hamming": torch.hamming_window,
"blackman": torch.blackman_window,
"bartlett": torch.bartlett_window,
"ones": torch.ones,
None: torch.ones,
}
# Ensure we can look up the window function
if window not in self.torch_windows:
raise ValueError(f"Got window value '{window}' but expected a member of {self.torch_windows.keys()}")
self.win_length = n_window_size
self.hop_length = n_window_stride
self._sample_rate = sample_rate
self._normalize_strategy = normalize
self._use_log = log
self._preemphasis_value = preemph
self.log_zero_guard_type = log_zero_guard_type
self.log_zero_guard_value: Union[str, float] = log_zero_guard_value
self.dither = dither
self.pad_to = pad_to
self.pad_value = pad_value
self.n_fft = n_fft
self._mel_spec_extractor: torchaudio.transforms.MelSpectrogram = torchaudio.transforms.MelSpectrogram(
sample_rate=self._sample_rate,
win_length=self.win_length,
hop_length=self.hop_length,
n_mels=nfilt,
window_fn=self.torch_windows[window],
mel_scale="slaney",
norm=mel_norm,
n_fft=n_fft,
f_max=highfreq,
f_min=lowfreq,
wkwargs={"periodic": False},
)
def _resolve_log_zero_guard_value(self, dtype: torch.dtype) -> float:
if isinstance(self.log_zero_guard_value, float):
return self.log_zero_guard_value
return getattr(torch.finfo(dtype), self.log_zero_guard_value)
def _apply_dithering(self, signals: torch.Tensor) -> torch.Tensor:
if self.training and self.dither > 0.0:
noise = torch.randn_like(signals) * self.dither
signals = signals + noise
return signals
def _apply_preemphasis(self, signals: torch.Tensor) -> torch.Tensor:
if self._preemphasis_value is not None:
padded = torch.nn.functional.pad(signals, (1, 0))
signals = signals - self._preemphasis_value * padded[:, :-1]
return signals
def _compute_output_lengths(self, input_lengths: torch.Tensor) -> torch.Tensor:
out_lengths = input_lengths.div(self.hop_length, rounding_mode="floor").add(1).long()
return out_lengths
def _apply_pad_to(self, features: torch.Tensor) -> torch.Tensor:
# Only apply during training; else need to capture dynamic shape for exported models
if not self.training or self.pad_to == 0 or features.shape[-1] % self.pad_to == 0:
return features
pad_length = self.pad_to - (features.shape[-1] % self.pad_to)
return torch.nn.functional.pad(features, pad=(0, pad_length), value=self.pad_value)
def _apply_log(self, features: torch.Tensor) -> torch.Tensor:
if self._use_log:
zero_guard = self._resolve_log_zero_guard_value(features.dtype)
if self.log_zero_guard_type == "add":
features = features + zero_guard
elif self.log_zero_guard_type == "clamp":
features = features.clamp(min=zero_guard)
else:
raise ValueError(f"Unsupported log zero guard type: '{self.log_zero_guard_type}'")
features = features.log()
return features
def _extract_spectrograms(self, signals: torch.Tensor) -> torch.Tensor:
# Complex FFT needs to be done in single precision
with torch.amp.autocast("cuda", enabled=False):
features = self._mel_spec_extractor(waveform=signals)
return features
def _apply_normalization(self, features: torch.Tensor, lengths: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
# For consistency, this function always does a masked fill even if not normalizing.
mask: torch.Tensor = make_seq_mask_like(lengths=lengths, like=features, time_dim=-1, valid_ones=False)
features = features.masked_fill(mask, 0.0)
# Maybe don't normalize
if self._normalize_strategy is None:
return features
# Use the log zero guard for the sqrt zero guard
guard_value = self._resolve_log_zero_guard_value(features.dtype)
if self._normalize_strategy == "per_feature" or self._normalize_strategy == "all_features":
# 'all_features' reduces over each sample; 'per_feature' reduces over each channel
reduce_dim = 2
if self._normalize_strategy == "all_features":
reduce_dim = [1, 2]
# [B, D, T] -> [B, D, 1] or [B, 1, 1]
means = features.sum(dim=reduce_dim, keepdim=True).div(lengths.view(-1, 1, 1))
stds = (
features.sub(means)
.masked_fill(mask, 0.0)
.pow(2.0)
.sum(dim=reduce_dim, keepdim=True) # [B, D, T] -> [B, D, 1] or [B, 1, 1]
.div(lengths.view(-1, 1, 1) - 1) # assume biased estimator
.clamp(min=guard_value) # avoid sqrt(0)
.sqrt()
)
features = (features - means) / (stds + eps)
else:
# Deprecating constant std/mean
raise ValueError(f"Unsupported norm type: '{self._normalize_strategy}")
features = features.masked_fill(mask, 0.0)
return features
def forward(self, input_signal: torch.Tensor, length: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
feature_lengths = self._compute_output_lengths(input_lengths=length)
signals = self._apply_dithering(signals=input_signal)
signals = self._apply_preemphasis(signals=signals)
features = self._extract_spectrograms(signals=signals)
features = self._apply_log(features=features)
features = self._apply_normalization(features=features, lengths=feature_lengths)
features = self._apply_pad_to(features=features)
return features, feature_lengths
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/multi_head_attention.py
+1133
View File
File diff suppressed because it is too large. Load diff
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/regularization_utils.py
+67
View File
@@ -0,0 +1,67 @@
# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/utils/regularization_utils.py; training-only code removed.
from typing import List
def compute_stochastic_depth_drop_probs(
num_layers: int,
stochastic_depth_drop_prob: float = 0.0,
stochastic_depth_mode: str = "linear",
stochastic_depth_start_layer: int = 1,
) -> List[float]:
"""Computes drop probabilities for stochastic depth regularization technique.
The first layer is never dropped and the starting layer needs to be greater
or equal to 1.
Args:
num_layers (int): number of layers in the network.
stochastic_depth_drop_prob (float): if non-zero, will randomly drop
layers during training. The higher this value, the more often layers
are dropped. Defaults to 0.0.
stochastic_depth_mode (str): can be either "linear" or "uniform". If
set to "uniform", all layers have the same probability of drop. If
set to "linear", the drop probability grows linearly from 0 for the
first layer to the desired value for the final layer. Defaults to
"linear".
stochastic_depth_start_layer (int): starting layer for stochastic depth.
All layers before this will never be dropped. Note that drop
probability will be adjusted accordingly if mode is "linear" when
start layer is > 1. Defaults to 1.
Returns:
List[float]: list of drop probabilities for all layers
"""
if not (0 <= stochastic_depth_drop_prob < 1.0):
raise ValueError("stochastic_depth_drop_prob has to be in [0, 1).")
if not (1 <= stochastic_depth_start_layer <= num_layers):
raise ValueError("stochastic_depth_start_layer has to be in [1, num layers].")
# Layers before `stochastic_depth_start_layer` are never dropped
layer_drop_probs = [0.0] * stochastic_depth_start_layer
# Layers starting with `stochastic_depth_start_layer` may be dropped
if (L := num_layers - stochastic_depth_start_layer) > 0:
if stochastic_depth_mode == "linear":
# we start with 1/L * drop_prob and and end with the desired drop probability.
layer_drop_probs += [l / L * stochastic_depth_drop_prob for l in range(1, L + 1)]
elif stochastic_depth_mode == "uniform":
layer_drop_probs += [stochastic_depth_drop_prob] * L
else:
raise ValueError(
f'stochastic_depth_mode has to be one of ["linear", "uniform"]. Current value: {stochastic_depth_mode}'
)
return layer_drop_probs
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/subsampling.py
+628
View File
@@ -0,0 +1,628 @@
# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/asr/parts/submodules/subsampling.py; training-only code removed.
import math
import torch
import torch.nn as nn
from torch.nn import LayerNorm
from .causal_convs import CausalConv1D, CausalConv2D
from ..compat import logging
class StackingSubsampling(torch.nn.Module):
"""Stacking subsampling which simply stacks consecutive frames to reduce the sampling rate
Args:
subsampling_factor (int): The subsampling factor
feat_in (int): size of the input features
feat_out (int): size of the output features
norm (bool): whether to use an MLP layer after the stacking along with normalization. default is False.
"""
def __init__(self, subsampling_factor, feat_in, feat_out, norm=False):
super(StackingSubsampling, self).__init__()
self.subsampling_factor = subsampling_factor
self.proj_out = torch.nn.Linear(subsampling_factor * feat_in, feat_out)
if norm:
self.pre_norm = LayerNorm(feat_in)
else:
self.pre_norm = None
def get_sampling_frames(self):
return self.subsampling_factor
def get_streaming_cache_size(self):
return 0
def forward(self, x, lengths):
b, t, h = x.size()
pad_size = (self.subsampling_factor - (t % self.subsampling_factor)) % self.subsampling_factor
x = torch.nn.functional.pad(x, (0, 0, 0, pad_size))
if self.pre_norm is not None:
x = self.pre_norm(x)
_, t, _ = x.size()
x = torch.reshape(x, (b, t // self.subsampling_factor, h * self.subsampling_factor))
x = self.proj_out(x)
lengths = torch.div(lengths + pad_size, self.subsampling_factor, rounding_mode="floor")
return x, lengths
class ConvSubsampling(torch.nn.Module):
"""Convolutional subsampling which supports VGGNet and striding approach introduced in:
VGGNet Subsampling: Transformer-transducer: end-to-end speech recognition with self-attention (https://arxiv.org/pdf/1910.12977.pdf)
Striding Subsampling: "Speech-Transformer: A No-Recurrence Sequence-to-Sequence Model for Speech Recognition" by Linhao Dong et al. (https://ieeexplore.ieee.org/document/8462506)
Args:
subsampling (str): The subsampling technique from {"vggnet", "striding", "dw-striding"}
subsampling_factor (int): The subsampling factor which should be a power of 2
subsampling_conv_chunking_factor (int): Input chunking factor which can be -1 (no chunking)
1 (auto) or a power of 2. Default is 1
feat_in (int): size of the input features
feat_out (int): size of the output features
conv_channels (int): Number of channels for the convolution layers.
activation (Module): activation function, default is nn.ReLU()
"""
def __init__(
self,
subsampling,
subsampling_factor,
feat_in,
feat_out,
conv_channels,
subsampling_conv_chunking_factor=1,
activation=nn.ReLU(),
is_causal=False,
):
super(ConvSubsampling, self).__init__()
self._subsampling = subsampling
self._conv_channels = conv_channels
self._feat_in = feat_in
self._feat_out = feat_out
if subsampling_factor % 2 != 0:
raise ValueError("Sampling factor should be a multiply of 2!")
self._sampling_num = int(math.log(subsampling_factor, 2))
self.subsampling_factor = subsampling_factor
self.is_causal = is_causal
if (
subsampling_conv_chunking_factor != -1
and subsampling_conv_chunking_factor != 1
and subsampling_conv_chunking_factor % 2 != 0
):
raise ValueError("subsampling_conv_chunking_factor should be -1, 1, or a power of 2")
self.subsampling_conv_chunking_factor = subsampling_conv_chunking_factor
in_channels = 1
layers = []
if subsampling == "vggnet":
self._stride = 2
self._kernel_size = 2
self._ceil_mode = True
self._left_padding = 0
self._right_padding = 0
for i in range(self._sampling_num):
layers.append(
torch.nn.Conv2d(
in_channels=in_channels, out_channels=conv_channels, kernel_size=3, stride=1, padding=1
)
)
layers.append(activation)
layers.append(
torch.nn.Conv2d(
in_channels=conv_channels, out_channels=conv_channels, kernel_size=3, stride=1, padding=1
)
)
layers.append(activation)
layers.append(
torch.nn.MaxPool2d(
kernel_size=self._kernel_size,
stride=self._stride,
padding=self._left_padding,
ceil_mode=self._ceil_mode,
)
)
in_channels = conv_channels
elif subsampling == "dw_striding":
self._stride = 2
self._kernel_size = 3
self._ceil_mode = False
if self.is_causal:
self._left_padding = self._kernel_size - 1
self._right_padding = self._stride - 1
self._max_cache_len = subsampling_factor + 1
else:
self._left_padding = (self._kernel_size - 1) // 2
self._right_padding = (self._kernel_size - 1) // 2
self._max_cache_len = 0
# Layer 1
if self.is_causal:
layers.append(
CausalConv2D(
in_channels=in_channels,
out_channels=conv_channels,
kernel_size=self._kernel_size,
stride=self._stride,
padding=None,
)
)
else:
layers.append(
torch.nn.Conv2d(
in_channels=in_channels,
out_channels=conv_channels,
kernel_size=self._kernel_size,
stride=self._stride,
padding=self._left_padding,
)
)
in_channels = conv_channels
layers.append(activation)
for i in range(self._sampling_num - 1):
if self.is_causal:
layers.append(
CausalConv2D(
in_channels=in_channels,
out_channels=in_channels,
kernel_size=self._kernel_size,
stride=self._stride,
padding=None,
groups=in_channels,
)
)
else:
layers.append(
torch.nn.Conv2d(
in_channels=in_channels,
out_channels=in_channels,
kernel_size=self._kernel_size,
stride=self._stride,
padding=self._left_padding,
groups=in_channels,
)
)
layers.append(
torch.nn.Conv2d(
in_channels=in_channels,
out_channels=conv_channels,
kernel_size=1,
stride=1,
padding=0,
groups=1,
)
)
layers.append(activation)
in_channels = conv_channels
elif subsampling == "striding":
self._stride = 2
self._kernel_size = 3
self._ceil_mode = False
if self.is_causal:
self._left_padding = self._kernel_size - 1
self._right_padding = self._stride - 1
self._max_cache_len = subsampling_factor + 1
else:
self._left_padding = (self._kernel_size - 1) // 2
self._right_padding = (self._kernel_size - 1) // 2
self._max_cache_len = 0
for i in range(self._sampling_num):
if self.is_causal:
layers.append(
CausalConv2D(
in_channels=in_channels,
out_channels=conv_channels,
kernel_size=self._kernel_size,
stride=self._stride,
padding=None,
)
)
else:
layers.append(
torch.nn.Conv2d(
in_channels=in_channels,
out_channels=conv_channels,
kernel_size=self._kernel_size,
stride=self._stride,
padding=self._left_padding,
)
)
layers.append(activation)
in_channels = conv_channels
elif subsampling == "striding_conv1d":
in_channels = feat_in
self._stride = 2
self._kernel_size = 5
self._ceil_mode = False
if self.is_causal:
self._left_padding = self._kernel_size - 1
self._right_padding = self._stride - 1
self._max_cache_len = subsampling_factor + 1
else:
self._left_padding = (self._kernel_size - 1) // 2
self._right_padding = (self._kernel_size - 1) // 2
self._max_cache_len = 0
for i in range(self._sampling_num):
if self.is_causal:
layers.append(
CausalConv1D(
in_channels=in_channels,
out_channels=feat_out if self._sampling_num == i + 1 else conv_channels,
kernel_size=self._kernel_size,
stride=self._stride,
padding=None,
)
)
else:
layers.append(
torch.nn.Conv1d(
in_channels=in_channels,
out_channels=feat_out if self._sampling_num == i + 1 else conv_channels,
kernel_size=self._kernel_size,
stride=self._stride,
padding=self._left_padding,
)
)
layers.append(activation)
in_channels = conv_channels
elif subsampling == "dw_striding_conv1d":
in_channels = feat_in
self._stride = 2
self._kernel_size = 5
self._ceil_mode = False
self._left_padding = (self._kernel_size - 1) // 2
self._right_padding = (self._kernel_size - 1) // 2
# Layer 1
layers.extend(
[
torch.nn.Conv1d(
in_channels=in_channels,
out_channels=in_channels,
kernel_size=self._kernel_size,
stride=self._stride,
padding=self._left_padding,
groups=in_channels,
),
torch.nn.Conv1d(
in_channels=in_channels,
out_channels=feat_out if self._sampling_num == 1 else conv_channels,
kernel_size=1,
stride=1,
padding=0,
groups=1,
),
]
)
in_channels = conv_channels
layers.append(activation)
for i in range(self._sampling_num - 1):
layers.extend(
[
torch.nn.Conv1d(
in_channels=in_channels,
out_channels=in_channels,
kernel_size=self._kernel_size,
stride=self._stride,
padding=self._left_padding,
groups=in_channels,
),
torch.nn.Conv1d(
in_channels=in_channels,
out_channels=feat_out if self._sampling_num == i + 2 else conv_channels,
kernel_size=1,
stride=1,
padding=0,
groups=1,
),
]
)
layers.append(activation)
in_channels = conv_channels
else:
raise ValueError(f"Not valid sub-sampling: {subsampling}!")
if subsampling in ["vggnet", "dw_striding", "striding"]:
in_length = torch.tensor(feat_in, dtype=torch.float)
out_length = calc_length(
lengths=in_length,
all_paddings=self._left_padding + self._right_padding,
kernel_size=self._kernel_size,
stride=self._stride,
ceil_mode=self._ceil_mode,
repeat_num=self._sampling_num,
)
self.out = torch.nn.Linear(conv_channels * int(out_length), feat_out)
self.conv2d_subsampling = True
elif subsampling in ["striding_conv1d", "dw_striding_conv1d"]:
self.out = None
self.conv2d_subsampling = False
else:
raise ValueError(f"Not valid sub-sampling: {subsampling}!")
self.conv = torch.nn.Sequential(*layers)
def get_sampling_frames(self):
return [1, self.subsampling_factor]
def get_streaming_cache_size(self):
return [0, self.subsampling_factor + 1]
def forward(self, x, lengths):
lengths = calc_length(
lengths,
all_paddings=self._left_padding + self._right_padding,
kernel_size=self._kernel_size,
stride=self._stride,
ceil_mode=self._ceil_mode,
repeat_num=self._sampling_num,
)
# Unsqueeze Channel Axis
if self.conv2d_subsampling:
x = x.unsqueeze(1)
# Transpose to Channel First mode
else:
x = x.transpose(1, 2)
# split inputs if chunking_factor is set
if self.subsampling_conv_chunking_factor != -1 and self.conv2d_subsampling:
if self.subsampling_conv_chunking_factor == 1:
# if subsampling_conv_chunking_factor is 1, we split only if needed
# avoiding a bug / feature limiting indexing of tensors to 2**31
# see https://github.com/pytorch/pytorch/issues/80020
x_ceil = 2**31 / self._conv_channels * self._stride * self._stride
if torch.numel(x) > x_ceil:
need_to_split = True
else:
need_to_split = False
else:
# if subsampling_conv_chunking_factor > 1 we always split
need_to_split = True
if need_to_split:
x, success = self.conv_split_by_batch(x)
if not success: # if unable to split by batch, try by channel
if self._subsampling == "dw_striding":
x = self.conv_split_by_channel(x)
else:
x = self.conv(x) # try anyway
else:
x = self.conv(x)
else:
x = self.conv(x)
# Flatten Channel and Frequency Axes
if self.conv2d_subsampling:
b, c, t, f = x.size()
x = self.out(x.transpose(1, 2).reshape(b, t, -1))
# Transpose to Channel Last mode
else:
x = x.transpose(1, 2)
return x, lengths
def reset_parameters(self):
# initialize weights
if self._subsampling == "dw_striding":
with torch.no_grad():
# init conv
scale = 1.0 / self._kernel_size
dw_max = (self._kernel_size**2) ** -0.5
pw_max = self._conv_channels**-0.5
torch.nn.init.uniform_(self.conv[0].weight, -scale, scale)
torch.nn.init.uniform_(self.conv[0].bias, -scale, scale)
for idx in range(2, len(self.conv), 3):
torch.nn.init.uniform_(self.conv[idx].weight, -dw_max, dw_max)
torch.nn.init.uniform_(self.conv[idx].bias, -dw_max, dw_max)
torch.nn.init.uniform_(self.conv[idx + 1].weight, -pw_max, pw_max)
torch.nn.init.uniform_(self.conv[idx + 1].bias, -pw_max, pw_max)
# init fc (80 * 64 = 5120 from https://github.com/kssteven418/Squeezeformer/blob/13c97d6cf92f2844d2cb3142b4c5bfa9ad1a8951/src/models/conformer_encoder.py#L487
fc_scale = (self._feat_out * self._feat_in / self._sampling_num) ** -0.5
torch.nn.init.uniform_(self.out.weight, -fc_scale, fc_scale)
torch.nn.init.uniform_(self.out.bias, -fc_scale, fc_scale)
def conv_split_by_batch(self, x):
"""Tries to split input by batch, run conv and concat results"""
b, _, _, _ = x.size()
if b == 1: # can't split if batch size is 1
return x, False
if self.subsampling_conv_chunking_factor > 1:
cf = self.subsampling_conv_chunking_factor
logging.debug(f"using manually set chunking factor: {cf}")
else:
# avoiding a bug / feature limiting indexing of tensors to 2**31
# see https://github.com/pytorch/pytorch/issues/80020
x_ceil = 2**31 / self._conv_channels * self._stride * self._stride
p = math.ceil(math.log(torch.numel(x) / x_ceil, 2))
cf = 2**p
logging.debug(f"using auto set chunking factor: {cf}")
new_batch_size = b // cf
if new_batch_size == 0: # input is too big
return x, False
logging.debug(f"conv subsampling: using split batch size {new_batch_size}")
return torch.cat([self.conv(chunk) for chunk in torch.split(x, new_batch_size, 0)]), True
def conv_split_by_channel(self, x):
"""For dw convs, tries to split input by time, run conv and concat results"""
x = self.conv[0](x) # full conv2D
x = self.conv[1](x) # activation
for i in range(self._sampling_num - 1):
_, c, t, _ = x.size()
if self.subsampling_conv_chunking_factor > 1:
cf = self.subsampling_conv_chunking_factor
logging.debug(f"using manually set chunking factor: {cf}")
else:
# avoiding a bug / feature limiting indexing of tensors to 2**31
# see https://github.com/pytorch/pytorch/issues/80020
p = math.ceil(math.log(torch.numel(x) / 2**31, 2))
cf = 2**p
logging.debug(f"using auto set chunking factor: {cf}")
new_c = int(c // cf)
if new_c == 0:
logging.warning(f"chunking factor {cf} is too high; splitting down to one channel.")
new_c = 1
new_t = int(t // cf)
if new_t == 0:
logging.warning(f"chunking factor {cf} is too high; splitting down to one timestep.")
new_t = 1
logging.debug(f"conv dw subsampling: using split C size {new_c} and split T size {new_t}")
x = self.channel_chunked_conv(self.conv[i * 3 + 2], new_c, x) # conv2D, depthwise
# splitting pointwise convs by time
x = torch.cat([self.conv[i * 3 + 3](chunk) for chunk in torch.split(x, new_t, 2)], 2) # conv2D, pointwise
x = self.conv[i * 3 + 4](x) # activation
return x
def channel_chunked_conv(self, conv, chunk_size, x):
"""Performs channel chunked convolution"""
ind = 0
out_chunks = []
for chunk in torch.split(x, chunk_size, 1):
step = chunk.size()[1]
if self.is_causal:
chunk = nn.functional.pad(
chunk, pad=(self._kernel_size - 1, self._stride - 1, self._kernel_size - 1, self._stride - 1)
)
ch_out = nn.functional.conv2d(
chunk,
conv.weight[ind : ind + step, :, :, :],
bias=conv.bias[ind : ind + step],
stride=self._stride,
padding=0,
groups=step,
)
else:
ch_out = nn.functional.conv2d(
chunk,
conv.weight[ind : ind + step, :, :, :],
bias=conv.bias[ind : ind + step],
stride=self._stride,
padding=self._left_padding,
groups=step,
)
out_chunks.append(ch_out)
ind += step
return torch.cat(out_chunks, 1)
def change_subsampling_conv_chunking_factor(self, subsampling_conv_chunking_factor: int):
if (
subsampling_conv_chunking_factor != -1
and subsampling_conv_chunking_factor != 1
and subsampling_conv_chunking_factor % 2 != 0
):
raise ValueError("subsampling_conv_chunking_factor should be -1, 1, or a power of 2")
self.subsampling_conv_chunking_factor = subsampling_conv_chunking_factor
def calc_length(lengths, all_paddings, kernel_size, stride, ceil_mode, repeat_num=1):
"""Calculates the output length of a Tensor passed through a convolution or max pooling layer"""
add_pad: float = all_paddings - kernel_size
one: float = 1.0
for i in range(repeat_num):
lengths = torch.div(lengths.to(dtype=torch.float) + add_pad, stride) + one
if ceil_mode:
lengths = torch.ceil(lengths)
else:
lengths = torch.floor(lengths)
return lengths.to(dtype=torch.int)
class SubsamplingReductionModule(nn.Module):
"""Downsamples the audio signal in time dimension."""
def __init__(self, reduction: str, d_model: int, reduction_factor: int = 2):
super().__init__()
assert reduction in ["pooling", "striding"]
self.reduction = reduction
self.d_model = d_model
self._sampling_num = int(math.log(reduction_factor, 2))
if reduction == "pooling":
self.reduction_enc = nn.MaxPool1d(kernel_size=reduction_factor)
self.padding = 0
self.kernel_size = self.reduction_enc.kernel_size
self.stride = self.reduction_enc.stride
elif reduction == "striding":
self.reduction_enc = ConvSubsampling(
subsampling="striding",
subsampling_factor=reduction_factor,
feat_in=d_model,
feat_out=d_model,
conv_channels=d_model,
activation=nn.ReLU(),
is_causal=False,
)
def forward(self, x, lengths):
"""Shapes:
- x: [B, T, C]
- lengths: [B]
"""
if self.reduction == "striding":
x, lengths = self.reduction_enc(x=x, lengths=lengths)
else:
x = torch.transpose(x, 1, 2) # [B, C, T]
lengths = calc_length(
lengths=lengths,
all_paddings=self.padding,
kernel_size=self.kernel_size,
stride=self.stride,
ceil_mode=False,
repeat_num=self._sampling_num,
)
x = self.reduction_enc(x)
x = torch.transpose(x, 1, 2) # [B, T, C]
return x, lengths
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/compat.py
+644
View File
@@ -0,0 +1,644 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat); compatibility shims
# replacing nemo.core / nemo.utils machinery for dependency-free inference.
# Each vendored function below is an exact copy from its annotated original source path;
# only imports were adjusted. Stub classes replace NeMo base classes whose functionality
# is not needed at inference time.
import importlib
import logging as _stdlib_logging
import math
import warnings
from contextlib import contextmanager, nullcontext
from dataclasses import dataclass
from typing import Dict, Optional
import torch
import torch.nn as nn
from omegaconf import OmegaConf
from safetensors.torch import load_file
from transformers import AutoConfig, AutoModelForCausalLM
# ==============================================================================
# Logging shim (replaces the nemo.utils logging import)
# ==============================================================================
class logging_mode:
"""Stub of nemo.utils.logging's dedup modes (used as `mode=logging_mode.ONCE`)."""
EACH = 0
ONCE = 1
class _NemoLoggerShim:
"""stdlib-logging based replacement for the nemo.utils logging module.
NeMo's logger accepts an extra `mode=` kwarg (deduplication); it is accepted
and ignored here, with a simple set-based dedup for `mode=logging_mode.ONCE`.
"""
def __init__(self, logger):
self._logger = logger
self._once_seen = set()
def _log(self, level, msg, args, mode=None):
if mode == logging_mode.ONCE:
key = (level, str(msg))
if key in self._once_seen:
return
self._once_seen.add(key)
self._logger.log(level, msg, *args)
def debug(self, msg, *args, mode=None, **kwargs):
self._log(_stdlib_logging.DEBUG, msg, args, mode=mode)
def info(self, msg, *args, mode=None, **kwargs):
self._log(_stdlib_logging.INFO, msg, args, mode=mode)
def warning(self, msg, *args, mode=None, **kwargs):
self._log(_stdlib_logging.WARNING, msg, args, mode=mode)
def error(self, msg, *args, mode=None, **kwargs):
self._log(_stdlib_logging.ERROR, msg, args, mode=mode)
def critical(self, msg, *args, mode=None, **kwargs):
self._log(_stdlib_logging.CRITICAL, msg, args, mode=mode)
logging = _NemoLoggerShim(_stdlib_logging.getLogger("nemo_vendored"))
# ==============================================================================
# typecheck no-op (replaces nemo.core.classes.common.typecheck)
# ==============================================================================
class typecheck:
"""No-op replacement for NeMo's typecheck decorator.
Supports the two usages found in the vendored files:
@typecheck()
def forward(...)
@typecheck.disable_checks()
def forward(...)
with typecheck.disable_checks():
...
"""
def __init__(self, *args, **kwargs):
pass
def __call__(self, func):
return func
@staticmethod
@contextmanager
def disable_checks():
yield
# ==============================================================================
# Neural type stubs (replace nemo.core.neural_types)
# ==============================================================================
# ==============================================================================
# from_config_dict target resolver (replaces hydra-based NeMo Serialization)
# ==============================================================================
# Maps the `_target_` strings that appear in the NVIDIA-NemotronLabs-VoiceChat-11B
# checkpoint config.json to the vendored classes. Extend as needed.
_TARGET_REGISTRY = {
"nemo.collections.asr.modules.ConformerEncoder": (".asr.conformer_encoder", "ConformerEncoder"),
"nemo.collections.asr.modules.conformer_encoder.ConformerEncoder": (
".asr.conformer_encoder",
"ConformerEncoder",
),
"nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor": (
".asr.audio_preprocessing",
"AudioToMelSpectrogramPreprocessor",
),
"nemo.collections.asr.modules.audio_preprocessing.AudioToMelSpectrogramPreprocessor": (
".asr.audio_preprocessing",
"AudioToMelSpectrogramPreprocessor",
),
"nemo.collections.speechlm2.modules.perception.IdentityConnector": (".perception", "IdentityConnector"),
"nemo.collections.speechlm2.modules.perception.AudioPerceptionModule": (".perception", "AudioPerceptionModule"),
}
_PACKAGE = __name__.rsplit(".", 1)[0] # ...nemotron_voicechat.nemo_vendored
def from_config_dict(config):
"""Instantiate a vendored class from an OmegaConf/dict config with a `_target_` key.
Hydra-free replacement for nemo.core.classes.common.Serialization.from_config_dict.
Like hydra.utils.instantiate with the default `_convert_="none"`, all keys except
`_target_` are passed as keyword arguments (nested containers stay OmegaConf nodes).
"""
cfg = OmegaConf.create(config) if not OmegaConf.is_config(config) else config
if "_target_" not in cfg:
raise ValueError(f"from_config_dict requires a '_target_' key in the config; got keys: {list(cfg.keys())}")
target = cfg["_target_"]
if target not in _TARGET_REGISTRY:
raise ValueError(
f"Unknown _target_ '{target}' — not in the vendored nemo_vendored.compat._TARGET_REGISTRY. "
f"Known targets: {sorted(_TARGET_REGISTRY)}. Add a mapping if this class was vendored."
)
module_name, class_name = _TARGET_REGISTRY[target]
cls = getattr(importlib.import_module(module_name, package=_PACKAGE), class_name)
kwargs = {str(k): cfg[k] for k in cfg if k != "_target_"}
return cls(**kwargs)
# ==============================================================================
# NeuralModule / Exportable / mixin stubs (replace nemo.core.classes.*)
# ==============================================================================
class NeuralModule(nn.Module):
"""Slim replacement of nemo.core.classes.module.NeuralModule."""
@property
def num_weights(self):
num: int = 0
for p in self.parameters():
if p.requires_grad:
num += p.numel()
return num
@staticmethod
def from_config_dict(config):
return from_config_dict(config)
def unfreeze(self, partial: bool = False) -> None:
"""Unfreeze all parameters for training. (exact copy of NeuralModule.unfreeze)"""
if partial and not hasattr(self, "_frozen_grad_map"):
raise ValueError("Cannot unfreeze partially without first freezing the module with `freeze()`")
for pname, param in self.named_parameters():
if not partial:
param.requires_grad = True
elif pname in self._frozen_grad_map:
param.requires_grad = self._frozen_grad_map[pname]
if hasattr(self, "_frozen_grad_map"):
delattr(self, "_frozen_grad_map")
self.train()
class Exportable:
"""Stub of nemo.core.classes.exportable.Exportable (ONNX export machinery removed)."""
@property
def input_names(self):
return []
@property
def output_names(self):
return []
class AccessMixin:
"""Stub of nemo.core.classes.mixins.access_mixins.AccessMixin.
Tensor-capture registry is training/interctc-only; at inference access is
always disabled so guarded blocks are skipped.
"""
@property
def access_cfg(self):
return {}
def is_access_enabled(self, guid=None):
return False
def register_accessible_tensor(self, name, tensor):
pass
class AttentionAdapterModuleMixin:
"""Stub of nemo.collections.asr.parts.submodules.adapters.attention_adapter_mixin.
Adapters (PEFT) are training-only; report none available."""
def is_adapter_available(self) -> bool:
return False
class AdapterModuleMixin(AttentionAdapterModuleMixin):
"""Stub of nemo.core.classes.mixins.adapter_mixins.AdapterModuleMixin."""
pass
# ==============================================================================
# StreamingEncoder — exact copy of
# nemo/collections/asr/parts/mixins/streaming.py (imports adjusted only)
# ==============================================================================
from abc import ABC, abstractmethod # noqa: E402
class StreamingEncoder(ABC):
@abstractmethod
def setup_streaming_params(
self,
max_look_ahead: int = 10000,
):
"""
This function sets the needed values and parameters to perform streaming. The configuration (CacheAwareStreamingConfig) need to be stored in self.streaming_cfg.
The streaming configuration is needed to simulate streaming inference. It would set the following
"""
pass
# ==============================================================================
# CacheAwareStreamingConfig — exact copy of the dataclass from
# nemo/collections/asr/models/configs/asr_models_config.py
# ==============================================================================
@dataclass
class CacheAwareStreamingConfig:
chunk_size: int = 0 # the size of each chunk at each step, it can be a list of two integers to specify different chunk sizes for the first step and others
shift_size: int = 0 # the size of the shift in each step, it can be a list of two integers to specify different shift sizes for the first step and others
cache_drop_size: int = 0 # the number of steps to drop from the cache
last_channel_cache_size: int = 0 # the size of the needed cache for last channel layers
valid_out_len: int = (
0 # the number of the steps in the final output which are valid (have the same value as in the offline mode)
)
pre_encode_cache_size: int = 0 # the size of the needed cache for the pre-encoding part of the model to avoid caching inside the pre-encoding layers
drop_extra_pre_encoded: int = 0 # the number of steps to get dropped after the pre-encoding layer
last_channel_num: int = 0 # number of the last channel layers (like MHA layers) which need caching in the model
last_time_num: int = 0 # number of the last time layers (like convolutions) which need caching in the model
# ==============================================================================
# avoid_float16_autocast_context / avoid_bfloat16_autocast_context — exact
# copies of NeMo nemo/utils/cast_utils.py
# ==============================================================================
def avoid_float16_autocast_context():
"""
If the current autocast context is float16, cast it to bfloat16
if available (unless we're in jit) or float32
"""
if torch.is_autocast_enabled() and torch.get_autocast_gpu_dtype() == torch.float16:
if torch.jit.is_scripting() or torch.jit.is_tracing():
return torch.amp.autocast("cuda", dtype=torch.float32)
if torch.cuda.is_bf16_supported():
return torch.amp.autocast("cuda", dtype=torch.bfloat16)
else:
return torch.amp.autocast("cuda", dtype=torch.float32)
else:
return nullcontext()
# ==============================================================================
# fp32_precision — exact copy of nemo/collections/speechlm2/parts/precision.py
# ==============================================================================
@contextmanager
def fp32_precision():
"""
Workaround for precision related issues when training with bf16-true PyTorch Lightning precision setting.
In bf16-true, PTL changes PyTorch's default dtype, which may break implicit assumptions for some models.
This context manager restores default float32 precision and runs the computation in float32 autocast context.
"""
default_dtype = torch.get_default_dtype()
torch.set_default_dtype(torch.float32)
try:
with torch.amp.autocast(device_type="cuda" if torch.cuda.is_available() else "cpu", dtype=torch.float32):
yield
finally:
torch.set_default_dtype(default_dtype)
# ==============================================================================
# get_pad_id — exact copy of nemo/collections/speechlm2/data/utils.py
# ==============================================================================
def get_pad_id(tokenizer) -> int:
pad_id = tokenizer.pad
if pad_id is not None:
return pad_id
pad_id = tokenizer.unk_id
if pad_id is not None:
return pad_id
warnings.warn(
"The text tokenizer has no <pad> or <unk> tokens available, using ID 0 for padding (this may lead to silent bugs)."
)
return 0
# ==============================================================================
# resample — exact copies of resample/_get_sinc_resample_kernel/_apply_sinc_resample_kernel
# of nemo/collections/audio/parts/utils/resampling.py (itself vendored from
# torchaudio 2.6.0, BSD 2-Clause; see that file for the license text).
# nemo/collections/audio/parts/utils/transforms.py re-exports `resample` from there.
# ==============================================================================
def resample(
waveform: torch.Tensor,
orig_freq: int,
new_freq: int,
lowpass_filter_width: int = 6,
rolloff: float = 0.99,
resampling_method: str = "sinc_interp_hann",
beta: float | None = None,
) -> torch.Tensor:
r"""Resamples the waveform at the new frequency using bandlimited interpolation. :cite:`RESAMPLE`."""
if orig_freq <= 0.0 or new_freq <= 0.0:
raise ValueError("Original frequency and desired frequency should be positive")
if orig_freq == new_freq:
return waveform
gcd = math.gcd(int(orig_freq), int(new_freq))
kernel, width = _get_sinc_resample_kernel(
orig_freq,
new_freq,
gcd,
lowpass_filter_width,
rolloff,
resampling_method,
beta,
waveform.device,
waveform.dtype,
)
resampled = _apply_sinc_resample_kernel(waveform, orig_freq, new_freq, gcd, kernel, width)
return resampled
def _get_sinc_resample_kernel(
orig_freq: int,
new_freq: int,
gcd: int,
lowpass_filter_width: int = 6,
rolloff: float = 0.99,
resampling_method: str = "sinc_interp_hann",
beta: float | None = None,
device: torch.device = "cpu",
dtype: torch.dtype | None = None,
):
if not (int(orig_freq) == orig_freq and int(new_freq) == new_freq):
raise Exception(
"Frequencies must be of integer type to ensure quality resampling computation. "
"To work around this, manually convert both frequencies to integer values "
"that maintain their resampling rate ratio before passing them into the function. "
"Example: To downsample a 44100 hz waveform by a factor of 8, use "
"`orig_freq=8` and `new_freq=1` instead of `orig_freq=44100` and `new_freq=5512.5`. "
"For more information, please refer to https://github.com/pytorch/audio/issues/1487."
)
if resampling_method not in ["sinc_interp_hann", "sinc_interp_kaiser"]:
raise ValueError("Invalid resampling method: {}".format(resampling_method))
orig_freq = int(orig_freq) // gcd
new_freq = int(new_freq) // gcd
if lowpass_filter_width <= 0:
raise ValueError("Low pass filter width should be positive.")
base_freq = min(orig_freq, new_freq)
# This will perform antialiasing filtering by removing the highest frequencies.
# At first I thought I only needed this when downsampling, but when upsampling
# you will get edge artifacts without this, as the edge is equivalent to zero padding,
# which will add high freq artifacts.
base_freq *= rolloff
# The key idea of the algorithm is that x(t) can be exactly reconstructed from x[i] (tensor)
# using the sinc interpolation formula:
# x(t) = sum_i x[i] sinc(pi * orig_freq * (i / orig_freq - t))
# We can then sample the function x(t) with a different sample rate:
# y[j] = x(j / new_freq)
# or,
# y[j] = sum_i x[i] sinc(pi * orig_freq * (i / orig_freq - j / new_freq))
# We see here that y[j] is the convolution of x[i] with a specific filter, for which
# we take an FIR approximation, stopping when we see at least `lowpass_filter_width` zeros crossing.
# But y[j+1] is going to have a different set of weights and so on, until y[j + new_freq].
# Indeed:
# y[j + new_freq] = sum_i x[i] sinc(pi * orig_freq * ((i / orig_freq - (j + new_freq) / new_freq))
# = sum_i x[i] sinc(pi * orig_freq * ((i - orig_freq) / orig_freq - j / new_freq))
# = sum_i x[i + orig_freq] sinc(pi * orig_freq * (i / orig_freq - j / new_freq))
# so y[j+new_freq] uses the same filter as y[j], but on a shifted version of x by `orig_freq`.
# This will explain the F.conv1d after, with a stride of orig_freq.
width = math.ceil(lowpass_filter_width * orig_freq / base_freq)
# If orig_freq is still big after GCD reduction, most filters will be very unbalanced, i.e.,
# they will have a lot of almost zero values to the left or to the right...
# There is probably a way to evaluate those filters more efficiently, but this is kept for
# future work.
idx_dtype = dtype if dtype is not None else torch.float64
idx = torch.arange(-width, width + orig_freq, dtype=idx_dtype, device=device)[None, None] / orig_freq
t = torch.arange(0, -new_freq, -1, dtype=dtype, device=device)[:, None, None] / new_freq + idx
t *= base_freq
t = t.clamp_(-lowpass_filter_width, lowpass_filter_width)
# we do not use built in torch windows here as we need to evaluate the window
# at specific positions, not over a regular grid.
if resampling_method == "sinc_interp_hann":
window = torch.cos(t * math.pi / lowpass_filter_width / 2) ** 2
else:
# sinc_interp_kaiser
if beta is None:
beta = 14.769656459379492
beta_tensor = torch.tensor(float(beta))
window = torch.i0(beta_tensor * torch.sqrt(1 - (t / lowpass_filter_width) ** 2)) / torch.i0(beta_tensor)
t *= math.pi
scale = base_freq / orig_freq
kernels = torch.where(t == 0, torch.tensor(1.0).to(t), t.sin() / t)
kernels *= window * scale
if dtype is None:
kernels = kernels.to(dtype=torch.float32)
return kernels, width
def _apply_sinc_resample_kernel(
waveform: torch.Tensor,
orig_freq: int,
new_freq: int,
gcd: int,
kernel: torch.Tensor,
width: int,
):
if not waveform.is_floating_point():
raise TypeError(f"Expected floating point type for waveform tensor, but received {waveform.dtype}.")
orig_freq = int(orig_freq) // gcd
new_freq = int(new_freq) // gcd
# pack batch
shape = waveform.size()
waveform = waveform.view(-1, shape[-1])
num_wavs, length = waveform.shape
waveform = torch.nn.functional.pad(waveform, (width, width + orig_freq))
resampled = torch.nn.functional.conv1d(waveform[:, None], kernel, stride=orig_freq)
resampled = resampled.transpose(1, 2).reshape(num_wavs, -1)
target_length = torch.ceil(torch.as_tensor(new_freq * length / orig_freq)).long()
resampled = resampled[..., :target_length]
# unpack batch
resampled = resampled.view(shape[:-1] + resampled.shape[-1:])
return resampled
# ==============================================================================
# Swish — exact copy of nemo/collections/asr/parts/utils/activations.py
# ==============================================================================
class Swish(nn.SiLU):
"""
Swish activation function introduced in 'https://arxiv.org/abs/1710.05941'
Mathematically identical to SiLU. See note in nn.SiLU for references.
"""
# ==============================================================================
# activation_registry — exact copy of nemo/collections/common/parts/utils.py
# ==============================================================================
activation_registry = {
"identity": nn.Identity,
"hardtanh": nn.Hardtanh,
"relu": nn.ReLU,
"selu": nn.SELU,
"swish": nn.SiLU,
"silu": nn.SiLU,
"gelu": nn.GELU,
}
# ==============================================================================
# FusedBatchNorm1d — exact copy of nemo/collections/asr/parts/submodules/batchnorm.py
# ==============================================================================
class FusedBatchNorm1d(nn.Module):
"""
Fused BatchNorm to use in Conformer to improve accuracy in finetuning with TTS scenario
Drop-in replacement for BatchNorm1d with simple affine projection
"""
def __init__(self, num_features: int):
"""
Args:
num_features: number of channels, see original BatchNorm1d documentation
"""
super().__init__()
self.num_features = num_features
self.weight = nn.Parameter(torch.ones(num_features))
self.bias = nn.Parameter(torch.zeros(num_features))
def forward(self, x: torch.Tensor):
if x.dim() == 3:
return x * self.weight.unsqueeze(-1) + self.bias.unsqueeze(-1)
assert x.dim() == 2
return x * self.weight + self.bias
# ==============================================================================
# load_pretrained_hf / set_model_dict_for_partial_init / load_checkpoint —
# exact copies of nemo/collections/speechlm2/parts/pretrained.py
# (only the functions used by the kept inference paths were vendored)
# ==============================================================================
def load_pretrained_hf(
model_path_or_name: str, pretrained_weights: bool = True, dtype=torch.float32, trust_remote_code: bool = True
):
"""
Load pretrained HuggingFace AutoModelForCausalLM.
Setting ``pretrained_weights=False`` returns a model that has identical architecture with the checkpoint,
but is randomly initialized.
Args:
model_path_or_name: Path or name of the model to load
pretrained_weights: Whether to load pretrained weights (True) or random init (False)
dtype: Data type for the model
trust_remote_code: Whether to trust remote code when loading model (needed for some models like Nemotron)
"""
if pretrained_weights:
return AutoModelForCausalLM.from_pretrained(
model_path_or_name, torch_dtype=dtype, trust_remote_code=trust_remote_code
)
else:
config = AutoConfig.from_pretrained(model_path_or_name, trust_remote_code=trust_remote_code)
return AutoModelForCausalLM.from_config(config, torch_dtype=dtype, trust_remote_code=trust_remote_code)
def set_model_dict_for_partial_init(
pretrained_dict: Dict[str, torch.Tensor], model_dict: Dict[str, torch.Tensor]
) -> Dict[str, torch.Tensor]:
"""
Partially initialize a model's state dictionary with a pretrained state dictionary.
This function safely copies compatible layers from a pretrained model into a new model,
ignoring layers with mismatched shapes or missing keys.
Steps:
1. Remove layers from the pretrained dictionary if their shape does not match the target model.
2. Keep only keys that exist in the target model.
3. Update the model dictionary with the filtered pretrained weights.
Args:
pretrained_dict (Dict[str, torch.Tensor]):
The state dictionary of the pretrained model.
model_dict (Dict[str, torch.Tensor]):
The state dictionary of the target model to be partially initialized.
Returns:
Dict[str, torch.Tensor]:
The updated model state dictionary with compatible layers loaded from the pretrained dictionary.
"""
# 1. Remove layers where pretrained shape differs from model shape
for k, v in list(pretrained_dict.items()):
if k in model_dict and hasattr(model_dict[k], "numel") and v.numel() != model_dict[k].numel():
del pretrained_dict[k]
logging.info(f" | > Layer with shape mismatch in the model definition: {k}")
# 2. Keep only keys that exist in the target model
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 3. Update model dictionary with filtered pretrained layers
model_dict.update(pretrained_dict)
logging.info(f" | > {len(pretrained_dict)} / {len(model_dict)} layers are restored.")
return model_dict
def load_checkpoint(checkpoint_path):
"""
Load a model checkpoint from disk.
Supports loading checkpoints stored in either PyTorch (`.ckpt`, `.pt`) or
SafeTensors (`.safetensors`) formats. All parameters are loaded onto CPU
regardless of the original device.
Args:
checkpoint_path (str):
Path to the checkpoint file. If the filename contains `.safetensors`,
it is loaded using the SafeTensors backend; otherwise, it is assumed
to be a PyTorch checkpoint containing a `state_dict` field.
Returns:
dict:
A state dictionary mapping parameter names to tensors.
"""
if ".safetensors" in checkpoint_path:
checkpoint_state = load_file(checkpoint_path, device="cpu")
else:
checkpoint_state = torch.load(checkpoint_path, map_location="cpu")["state_dict"]
return checkpoint_state
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/duplex_ear_tts.py
+1038
View File
File diff suppressed because it is too large. Load diff
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_commons.py
+260
View File
@@ -0,0 +1,260 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/speechlm2/modules/ear_tts_commons.py; training-only code removed.
import argparse
import glob
import importlib.machinery
import json
import os
import re
import shutil
import subprocess
import sys
from collections.abc import Mapping
from typing import Any
from omegaconf import DictConfig
from safetensors import safe_open
from torch import nn
from .compat import logging
# ==============================================================================
# Constants
# ==============================================================================
PYTHON_CONFIG_GETTER_NAME = "get_config"
CHECKPOINT_FORMAT = "checkpoint_{}/ema.safetensors"
CONFIG_NAME = "config.json"
GIT_HASH_NAME = "githash"
SCRIPT_PLACEHOLDER = "[[[<<<SCRIPT_PLACEHOLDER>>>]]]"
# ==============================================================================
# Configuration Class and Utilities
# ==============================================================================
def get_config_from_file(config_path: str) -> DictConfig:
"""
Loads a configuration from a JSON or Python file.
- For JSON files (`*.json`), it parses the file directly.
- For Python files (`*.py`), it imports the file as a module and calls a
`get_config()` function within it.
- It also supports a special syntax `path/to/config.py:config_name` to select
a specific configuration from a Python file that returns a dictionary of configs.
Args:
config_path (str): The path to the configuration file.
Returns:
Config: The loaded configuration object.
Raises:
AssertionError: If the file path is invalid, does not exist, or is not in
the expected format.
"""
match = re.search(r".+\.((json)|(py)|(py:.+))$", config_path)
assert match, f"Only Python (*.py) or JSON (*.json) files are supported, but got {config_path}."
py_config_name: str | None = None
if not (config_path.endswith(".py") or config_path.endswith(".json")):
config_path_split = config_path.split(":")
config_path = ":".join(config_path_split[:-1])
py_config_name = config_path_split[-1]
assert os.path.isfile(config_path), f"Configuration file not found at: {config_path}"
if config_path.endswith(".json"):
with open(config_path) as f:
config = json.load(f)
else:
config_module = importlib.machinery.SourceFileLoader("_config", config_path).load_module()
assert hasattr(config_module, PYTHON_CONFIG_GETTER_NAME), (
f"Python config file must define a `{PYTHON_CONFIG_GETTER_NAME}` function."
)
config = getattr(config_module, PYTHON_CONFIG_GETTER_NAME)(py_config_name)
assert isinstance(config, Mapping), f"`{PYTHON_CONFIG_GETTER_NAME}` must return a dictionary-like object."
cfg = DictConfig(config)
return cfg
def get_config() -> DictConfig:
"""
Parses command-line arguments to load the main configuration for a training run.
This function implements a hierarchical configuration loading strategy:
1. It checks if a `config.json` exists in the specified `--workdir`. If so, it loads it.
2. If a `--config` argument is also provided, it uses that file to update the
configuration loaded from the work directory.
3. If no config exists in the work directory, it requires the `--config` argument
to be provided as the base configuration.
This allows for resuming training from a work directory while also being able to
override specific parameters for a new run.
Returns:
Config: The final, consolidated configuration object.
"""
parser = argparse.ArgumentParser(description="Load training configuration.")
parser.add_argument("-c", "--config", type=str, default=None, help="Path to a Python or JSON configuration file.")
parser.add_argument("-w", "--workdir", type=str, required=True, help="Work directory to save logs and checkpoints.")
args = parser.parse_args()
workdir_path = args.workdir
config_save_path = os.path.join(workdir_path, CONFIG_NAME)
if os.path.exists(config_save_path):
logging.info(f"Resuming from work directory. Loading configuration from {config_save_path}.")
cfg = get_config_from_file(config_save_path)
if args.config and args.config != config_save_path:
logging.info(f"Updating loaded configuration with parameters from {args.config}.")
override_cfg = get_config_from_file(args.config)
cfg.update(override_cfg)
else:
assert args.config is not None, "A configuration file must be specified via `-c` or `--config` for a new run."
logging.info(f"Starting a new run. Loading configuration from {args.config}.")
cfg = get_config_from_file(args.config)
cfg.workdir_path = workdir_path
return cfg
def get_config_from_dir(workdir_path: str) -> DictConfig:
"""
A simple utility to load the configuration directly from a work directory.
Args:
workdir_path (str): The path to the work directory containing a `config.json`.
Returns:
DictConfig: The loaded configuration object.
"""
config_save_path = os.path.join(workdir_path, CONFIG_NAME)
cfg = get_config_from_file(config_save_path)
cfg.workdir_path = workdir_path
return cfg
# ==============================================================================
# Base Model Classes
# ==============================================================================
class PreTrainedModel(nn.Module):
config_class = DictConfig
"""
A base class for models to handle loading from pretrained checkpoints.
This class provides a common interface for initializing a model and loading
weights from a saved checkpoint, following a pattern similar to libraries
like Hugging Face's Transformers.
Args:
config (DictConfig | dict[str, Any]): A configuration object containing model hyperparameters.
"""
def __init__(self, config: DictConfig | dict[str, Any], *args, **kwargs):
super().__init__()
self.config = config if isinstance(config, self.config_class) else self.config_class(config)
@classmethod
def from_pretrained(
cls,
pretrained_dir: str,
cfg: DictConfig | dict[str, Any] | None = None,
checkpoint_regex: str = "checkpoint_*/ema.safetensors",
strict: bool = False,
**model_kwargs,
) -> "PreTrainedModel":
"""
Loads a pretrained model from a directory.
This method first loads the configuration file from the specified directory,
initializes the model with this configuration, and then loads the weights
from the latest checkpoint file found in that directory.
Args:
cls (type): The model class to instantiate.
pretrained_dir (str): The directory containing the pretrained model
config and checkpoint files.
cfg (DictConfig | dict[str, Any] | None, optional): An optional config object to override
the loaded config. Defaults to None.
checkpoint_regex (str, optional): A regex pattern to find the checkpoint
file. Defaults to "checkpoint_*/ema.safetensors".
strict (bool, optional): Whether to strictly enforce that the keys in
the checkpoint match the keys of the model.
Defaults to False.
**model_kwargs: Additional keyword arguments to pass to the model's
constructor.
Returns:
PreTrainedModel: An instance of the model with loaded weights.
"""
pretrained_cfg = get_config_from_dir(pretrained_dir).model
if cfg is not None:
pretrained_cfg.update(cfg)
logging.info(f"The loaded config of the pretrained model is updated to: {pretrained_cfg}")
model = cls(
pretrained_cfg,
**model_kwargs,
)
model_state_dict = {}
with safe_open(latest_checkpoint_path(pretrained_dir, checkpoint_regex), framework="pt", device="cpu") as f:
for key in f.keys():
model_state_dict[key] = f.get_tensor(key)
model.load_state_dict(model_state_dict, strict=strict)
return model
# ==============================================================================
# IO and Checkpointing Utilities
# ==============================================================================
def latest_checkpoint_path(dir_path: str, regex: str | None = None) -> str:
"""
Finds the path of the latest checkpoint file or directory in a directory.
The latest checkpoint is determined by sorting the filenames alphanumerically
and picking the last one. This assumes a naming convention like `checkpoint_1000.pt`,
`checkpoint_2000.pt`, etc.
Args:
dir_path (str): The directory to search for checkpoints.
regex (str | None, optional): A glob pattern to match checkpoint files. If None,
a default pattern is used. Defaults to None.
Returns:
str: The full path to the latest checkpoint file.
Raises:
AssertionError: If no files matching the regex are found in the directory.
"""
if regex is None:
regex = CHECKPOINT_FORMAT.format("*")
f_list = glob.glob(os.path.join(dir_path, regex))
if not f_list:
raise FileNotFoundError(f"No checkpoint files or directories found in {dir_path} matching '{regex}'")
# Sort files based on the integer values in their names
f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))
latest_path = f_list[-1]
logging.info(f"Latest checkpoint '{os.path.relpath(latest_path, start=dir_path)}' found in '{dir_path}'.")
return latest_path
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py
+1620
View File
File diff suppressed because it is too large. Load diff
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_vae_codec.py
+874
View File
@@ -0,0 +1,874 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/speechlm2/modules/ear_tts_vae_codec.py; training-only code removed.
import functools
import math
from collections.abc import Callable
from contextlib import contextmanager
from typing import Any, Concatenate
import torch
from omegaconf import DictConfig
from torch import Tensor, nn
from torch.nn import functional as F
@contextmanager
def disable_tf32():
prev = torch.backends.cudnn.allow_tf32
torch.backends.cudnn.allow_tf32 = False
try:
yield
finally:
torch.backends.cudnn.allow_tf32 = prev
# ==============================================================================
# Utility Functions
# ==============================================================================
def zero_module(module: nn.Module) -> nn.Module:
"""
Zeros out the parameters of a PyTorch module in-place.
This is a utility function that iterates through all parameters of a given
`nn.Module` and sets their values to zero. This is often used for specific
initialization strategies, for example in diffusion models where some layers
are initialized to zero.
From: https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/guided_diffusion/nn.py#L68
Args:
module (nn.Module): The PyTorch module to be zeroed.
Returns:
nn.Module: The same module with its parameters zeroed.
"""
for p in module.parameters():
# p.detach().zero_() performs the operation in-place without tracking it in autograd
p.detach().zero_()
return module
def sequence_mask(lengths: Tensor, max_length: int | None = None) -> Tensor:
"""
Creates a boolean mask from a 1D tensor of sequence lengths.
This function is useful for masking out padding in sequences. Given a tensor
of lengths, it produces a 2D boolean tensor where `mask[i, j]` is `True` if
`j < lengths[i]` and `False` otherwise.
Example:
>>> lengths = torch.tensor([1, 3, 2])
>>> sequence_mask(lengths)
tensor([[ True, False, False],
[ True, True, True],
[ True, True, False]])
Args:
lengths (Tensor): A 1D tensor of integer lengths. Shape: `[batch_size]`.
max_length (int | None, optional): The maximum length of the mask. If None,
it is inferred from the maximum value
in `lengths`. Defaults to None.
Returns:
Tensor: The boolean mask. Shape: `[batch_size, max_length]`.
"""
if max_length is None:
# If max_length is not provided, use the longest sequence length in the batch.
max_length = int(lengths.max().item())
# Create a range tensor from 0 to max_length - 1
x = torch.arange(max_length, dtype=lengths.dtype, device=lengths.device)
# Compare each length with the range tensor to create the mask.
# `x.unsqueeze(0)` is `[1, max_length]`
# `lengths.unsqueeze(1)` is `[batch_size, 1]`
# Broadcasting takes care of the comparison.
return x.unsqueeze(0) < lengths.unsqueeze(1)
# ==============================================================================
# Signal Processing Functions
# ==============================================================================
def spectrogram(
wav: Tensor,
n_fft: int,
hop_length: int,
win_length: int,
window_fn: Callable[Concatenate[int, ...], Tensor] = torch.hann_window,
) -> Tensor:
"""
Computes the Short-Time Fourier Transform (STFT) of a waveform with manual padding.
This implementation manually applies zero padding before computing the STFT.
This is done to center the analysis window at the beginning of the signal
without using the `center=True` argument in `torch.stft`, giving more control.
Args:
wav (Tensor): The input audio waveform.
Shape: [batch_size?, time_steps], where batch_size? is an
optional batch dimension.
n_fft (int): The size of the FFT.
hop_length (int): The number of samples between adjacent STFT columns.
win_length (int): The size of the window function.
window_fn (function, optional): The window function to apply.
Defaults to torch.hann_window.
Returns:
Tensor: The complex-valued spectrogram.
Shape: [batch_size?, n_fft // 2 + 1, num_frames]
"""
# Calculate the padding required on the left and right sides to center the frames.
pad_size_l = (n_fft - hop_length) // 2
pad_size_r = (n_fft - hop_length) - pad_size_l
# Use a torch.autocast context to perform STFT in float32 for precision.
with torch.autocast(device_type=wav.device.type, enabled=False):
# Apply reflection padding to the waveform.
wav = F.pad(wav.float(), (pad_size_l, pad_size_r))
# Create the window tensor on the same device as the waveform.
window = window_fn(win_length, dtype=torch.float, device=wav.device)
# Compute the STFT.
# `center=False` because we have already manually padded the signal.
spec = torch.stft(
wav,
n_fft,
hop_length=hop_length,
win_length=win_length,
window=window,
center=False,
normalized=False,
onesided=True,
return_complex=True,
)
return spec
def spec_to_wav(
spec: Tensor,
n_fft: int,
hop_length: int,
win_length: int,
window_fn: Callable[Concatenate[int, ...], Tensor] = torch.hann_window,
constrain_value_range: bool = False,
) -> Tensor:
"""
Converts a spectrogram back into a waveform using the overlap-add method.
This function is an approximate inverse of the `spectrogram` function.
Args:
spec (Tensor): The input complex-valued spectrogram.
Shape: [batch_size?, dim, time_steps], where batch_size?
is an optional batch dimension.
n_fft (int): The size of the FFT used to create the spectrogram.
hop_length (int): The number of samples between frames in the original signal.
win_length (int): The size of the window function used in the original signal.
window_fn (function, optional): The window function used. Currently only
`torch.hann_window` is supported.
constrain_value_range (bool, optional): If True, constrains the IFFT values
to be within the range of the window.
This ensures that the output values
remain within the range of -1.0 to 1.0.
Defaults to False.
Returns:
Tensor: The reconstructed waveform.
Shape: [batch_size?, time_steps]
Raises:
ValueError: If a window function other than `torch.hann_window` is provided.
"""
with torch.autocast(device_type=spec.device.type, enabled=False):
if window_fn != torch.hann_window:
raise ValueError(f"`window_fn` should be 'torch.hann_window', but got '{window_fn}'.")
# Calculate padding and number of frames
pad = (win_length - hop_length) // 2
T = spec.size(-1)
window = window_fn(win_length, device=spec.device)
# 1. Inverse FFT
# Convert from frequency domain back to time domain for each frame.
ifft = torch.fft.irfft(spec, n=n_fft, dim=-2, norm="backward")
window_unsqz = window.unsqueeze(-1)
# 2. Optionally constrain values
if constrain_value_range:
ifft = torch.where(
ifft >= 0,
torch.minimum(ifft, window_unsqz),
torch.maximum(ifft, -window_unsqz),
)
# 3. Apply window to the IFFT result
ifft = ifft * window_unsqz
# 4. Overlap and Add
# Use `torch.nn.functional.fold` to perform the overlap-add operation efficiently.
# This reconstructs the continuous signal from the windowed frames.
output_size = (T - 1) * hop_length + win_length
wav = F.fold(
ifft,
output_size=(1, output_size),
kernel_size=(1, win_length),
stride=(1, hop_length),
)[..., 0, 0, pad:-pad]
# 5. Calculate the window envelope for normalization
# This is necessary to correct for the energy added by overlapping windows.
window_sq = window.square().expand(T, -1).transpose(0, 1)
window_envelope = F.fold(
window_sq,
output_size=(1, output_size),
kernel_size=(1, win_length),
stride=(1, hop_length),
).squeeze()[pad:-pad]
# 6. Normalize the waveform
# Divide by the window envelope to get the final reconstructed signal.
assert (window_envelope > 1e-11).all(), "Window envelope has zero values, cannot normalize."
wav = wav / window_envelope
return wav
# ==============================================================================
# Basic Modules
# ==============================================================================
class CausalConv1dCache:
"""
A cache for managing states in causal 1D convolutions.
This class is used during autoregressive inference to store and update the
tail of the input to a causal convolution, which is used as padding for the
next time step. This avoids re-computing the entire sequence at each step.
"""
def __init__(self) -> None:
self.cache: dict[int | str, Tensor] = {}
def __getitem__(self, layer_id: int | str) -> Tensor:
"""Retrieves the cached tensor for a given layer."""
return self.cache[layer_id]
def update(
self,
states: Tensor,
layer_id: int | str,
padding: int,
padding_value: int = 0,
flush: bool = False,
) -> Tensor:
"""
Updates the cache for a specific layer and returns the padded input.
Args:
states (Tensor): The new input tensor for the current time step.
layer_id (int | str): An identifier for the convolutional layer.
padding (int): The amount of left padding required by the convolution.
padding_value (int, optional): The value to use for initial padding. Defaults to 0.
flush (bool, optional): If True, the cache for this layer is deleted
after use. Defaults to False.
Returns:
Tensor: The input states concatenated with the cached padding.
"""
device = states.device
dtype = states.dtype
b, c, t = states.size()
if layer_id not in self.cache:
# Initialize cache with zero padding if it's the first time step
padding_tensor = torch.zeros((b, c, padding), dtype=dtype, device=device) + padding_value
else:
padding_tensor = self.cache[layer_id]
assert padding_tensor.size(2) == padding
# Concatenate the cached padding with the new states
padded_states = torch.cat([padding_tensor, states], dim=2)
# Update the cache with the tail of the new padded states
self.cache[layer_id] = padded_states[:, :, -padding:]
if flush:
del self.cache[layer_id]
return padded_states
class LayerNormNd(nn.Module):
"""
A LayerNorm module that works for N-dimensional inputs.
This implementation normalizes over the channel dimension (dim=1), which is
a common setup for convolutional networks.
Args:
channels (int): The number of channels of the input tensor.
eps (float, optional): A value added to the denominator for numerical
stability. Defaults to 1e-6.
elementwise_affine (bool, optional): If True, this module has learnable
affine parameters (weight and bias).
Defaults to True.
bias (bool, optional): If True, this module has a learnable bias.
Defaults to True.
"""
def __init__(self, channels: int, eps=1e-6, elementwise_affine: bool = True, bias: bool = True):
super().__init__()
self.channels = channels
self.eps = eps
self.weight = nn.Parameter(torch.ones((channels,)), requires_grad=elementwise_affine)
self.bias = nn.Parameter(torch.zeros((channels,)), requires_grad=elementwise_affine and bias)
def forward(self, x: Tensor) -> Tensor:
# Calculate mean and reciprocal standard deviation over the channel dimension
mean = x.mean(1, keepdim=True)
x_shift = x - mean
# Using rsqrt for potentially better performance
x_rstd = torch.rsqrt(x_shift.pow(2).mean(1, keepdim=True) + self.eps)
# Reshape weight and bias to be broadcastable with the input tensor
shape = [-1 if i == 1 else 1 for i in range(x.ndim)]
# Apply normalization and affine transformation
return (x_shift * x_rstd) * self.weight.view(shape) + self.bias.view(shape)
class ConvNeXt1d(nn.Module):
"""
A 1D ConvNeXt block adapted for causal convolutions on audio signals.
This block is a core component of modern convolutional architectures, featuring
a depthwise convolution, layer normalization, and pointwise convolutions to
expand and contract the channel dimension, similar to an inverted bottleneck.
Implementation adapted from: https://github.com/charactr-platform/vocos
Args:
dim (int): Number of input and output channels.
intermediate_dim (int): Dimensionality of the intermediate (expanded) layer.
kernel_size (int): The kernel size for the causal depthwise convolution.
identity_init (bool, optional): If True, the final pointwise convolution
is initialized to zero, making the block
an identity function at the start of training.
Defaults to False.
layer_idx (int, optional): An index for this layer, used for caching. Defaults to 0.
"""
def __init__(
self,
dim: int,
intermediate_dim: int,
kernel_size: int,
identity_init: bool = False,
layer_idx: int = 0,
):
super().__init__()
self.layer_idx = layer_idx
self.kernel_size = kernel_size
# Depthwise convolution
self.dwconv = nn.Conv1d(dim, dim, kernel_size=kernel_size, groups=dim)
self.norm = LayerNormNd(dim)
self.pwconv1 = nn.Conv1d(dim, intermediate_dim, 1) # Pointwise/1x1 conv for expansion
self.act = nn.GELU()
# Pointwise/1x1 conv for projection
if identity_init:
self.pwconv2 = zero_module(nn.Conv1d(intermediate_dim, dim, 1))
else:
self.pwconv2 = nn.Conv1d(intermediate_dim, dim, 1)
def forward(self, x: Tensor, cache: CausalConv1dCache | None = None, flush: bool = False) -> Tensor:
residual = x
# Apply causal padding, either through a cache or manually
if cache is not None:
x = cache.update(x, self.layer_idx, self.kernel_size - 1, flush=flush)
else:
x = F.pad(x, [self.kernel_size - 1, 0]) # Left padding for causality
# Main ConvNeXt path
x = self.dwconv(x)
x = self.norm(x)
x = self.pwconv1(x)
x = self.act(x)
x = self.pwconv2(x)
# Add residual connection
x = residual + x
return x
class PreTrainedEMAVariance(nn.Module):
"""
Exponential Moving Average of Variance
"""
def __init__(self, initial_value: float = 1.0):
super().__init__()
self.variance = nn.Parameter(
torch.tensor(initial_value),
requires_grad=False,
)
def forward(self) -> Tensor:
return self.variance
class PreTrainedProbabilisticVQ(nn.Module):
def __init__(
self,
channels: int,
num_mixtures: int,
depth: int = 1,
):
super().__init__()
self.channels = channels
self.num_mixtures = num_mixtures
self.depth = depth
self.mus_list = nn.ParameterList(
[
nn.Parameter(
F.normalize(torch.randn(num_mixtures, channels), p=2.0, dim=1) * ((depth - i) / depth),
requires_grad=False,
)
for i in range(depth)
]
)
self._variance_list = nn.ModuleList([PreTrainedEMAVariance() for _ in range(depth)])
def encode(self, z: Tensor, return_z_q: bool = False) -> list[Tensor] | tuple[list[Tensor], Tensor]:
r = z
ids_sel = []
for i in range(self.depth):
mus = self.mus_list[i]
idx_sel = self._dist_sq(r, mus).argmin(-1) # [b, ?, h], [v, h] -> [b, ?]
r = r - F.embedding(idx_sel, mus)
ids_sel.append(idx_sel)
if return_z_q:
return ids_sel, z - r
return ids_sel
def decode(self, ids_sel: list[Tensor]) -> Tensor:
z = torch.zeros((*ids_sel[0].size(), self.channels), device=ids_sel[0].device)
for i in range(len(ids_sel)):
mus = self.mus_list[i]
z = z + F.embedding(ids_sel[i], mus)
return z # [b, ?, h]
def _dist_sq(self, z: Tensor, mus: Tensor) -> Tensor:
"""
z: [b, ?, d?, h]
mus: [d?, v, h]
"""
return (
z.pow(2).sum(-1, keepdim=True) # [b, ?, d?, 1]
+ mus.pow(2).sum(-1) # [d?, v]
- 2 * (z.unsqueeze(-2) @ mus.transpose(-1, -2)).squeeze(-2) # [b, ?, d?, h] , [d?, h, v] -> [b, ?, d?, v]
)
class Wav2Latent(nn.Module):
"""
An encoder model that transforms a raw waveform into a latent representation.
This model first converts the waveform to a spectrogram, then processes it
through a series of ConvNeXt blocks and downsampling convolutional layers
to produce a compressed latent tensor.
Args:
latent_size (int): The number of channels in the final latent representation.
n_fft (int): The FFT size for the initial spectrogram transformation.
hop_length (int): The hop length for the STFT.
base_hidden_size (int): The base number of channels for the hidden layers.
channel_mult (tuple[int, ...]): A tuple of multipliers for the hidden
size at each stage of downsampling.
rates (tuple[int, ...]): A tuple of downsampling factors (strides) for
the convolutional layers.
num_blocks (int): The number of ConvNeXt blocks per stage.
kernel_size (int): The kernel size for the ConvNeXt blocks.
groups (int): The number of groups for the downsampling convolutions.
"""
def __init__(
self,
latent_size: int = 1024,
n_fft: int = 32,
hop_length: int = 8,
base_hidden_size: int = 384,
channel_mult: tuple[int, ...] = (1, 2, 4),
rates: tuple[int, ...] = (8, 8, 8),
num_blocks: int = 3,
kernel_size: int = 7,
groups: int = 1,
):
super().__init__()
self.n_fft = n_fft
self.hop_length = hop_length
# Initial projection from spectrogram to hidden size
layers: list[nn.Module] = [nn.Conv1d(n_fft + 2, base_hidden_size * channel_mult[0], 1, bias=False)]
# Downsampling stages
for i in range(len(channel_mult)):
ch_mult, rate = channel_mult[i], rates[i]
hidden_size = base_hidden_size * ch_mult
# Add ConvNeXt blocks for this stage
for j in range(num_blocks):
layers.append(ConvNeXt1d(hidden_size, hidden_size * 4, kernel_size, True, layer_idx=i * num_blocks + j))
# Add downsampling convolution
next_hidden_size = base_hidden_size * channel_mult[i + 1] if i < len(channel_mult) - 1 else latent_size
layers.append(
nn.Conv1d(hidden_size, next_hidden_size, kernel_size=rate, stride=rate, bias=False, groups=groups)
)
self.layers = nn.ModuleList(layers)
def forward(self, x: Tensor, cache=None, flush: bool = False) -> Tensor:
if cache is not None:
raise NotImplementedError("Caching is not implemented for the encoder.")
# Convert waveform to spectrogram (magnitude and phase)
with torch.autocast(device_type=x.device.type, enabled=False):
spec = spectrogram(x.squeeze(1), n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.n_fft)
# Split complex spectrogram into real and imaginary, then treat as magnitude and phase
mag, ph = torch.view_as_real(spec).chunk(2, dim=-1)
x = torch.cat([mag, ph], 1).squeeze(-1)
# Pass through the network
for layer in self.layers:
if isinstance(layer, ConvNeXt1d):
x = layer(x, cache=cache, flush=flush)
else:
x = layer(x)
# Transpose to [batch, time, channels] for compatibility with transformers
x = x.transpose(-1, -2)
return x
class Latent2Wav(nn.Module):
"""
A decoder (vocoder) model that transforms a latent representation back into a raw waveform.
This model processes a latent tensor through a series of ConvNeXt blocks and
upsampling transposed convolutional layers to produce a spectrogram, which is
then converted back to a waveform using an inverse STFT.
Args:
latent_size (int): The number of channels in the input latent representation.
n_fft (int): The FFT size for the final spectrogram reconstruction.
hop_length (int): The hop length for the ISTFT.
base_hidden_size (int): The base number of channels for the hidden layers.
channel_mult (tuple[int, ...]): A tuple of multipliers for the hidden
size at each stage of upsampling.
rates (tuple[int, ...]): A tuple of upsampling factors (strides) for
the transposed convolutional layers.
num_blocks (int): The number of ConvNeXt blocks per stage.
kernel_size (int): The kernel size for the ConvNeXt blocks.
groups (int): The number of groups for the upsampling convolutions.
"""
def __init__(
self,
latent_size: int = 1024,
n_fft: int = 32,
hop_length: int = 8,
base_hidden_size: int = 384,
channel_mult: tuple[int, ...] = (4, 2, 1),
rates: tuple[int, ...] = (8, 8, 8),
num_blocks: int = 3,
kernel_size: int = 7,
groups=1,
):
super().__init__()
self.n_fft = n_fft
self.hop_length = hop_length
self.spec_cache_idx = (len(channel_mult)) * num_blocks
layers: list[nn.Module] = []
# Upsampling stages
for i in range(len(channel_mult)):
ch_mult, rate = channel_mult[i], rates[i]
hidden_size = base_hidden_size * ch_mult
# Add upsampling transposed convolution
in_size = base_hidden_size * channel_mult[i - 1] if i != 0 else latent_size
layers.append(
nn.ConvTranspose1d(in_size, hidden_size, kernel_size=rate, stride=rate, bias=False, groups=groups)
)
# Add ConvNeXt blocks for this stage
for j in range(num_blocks):
layers.append(ConvNeXt1d(hidden_size, hidden_size * 4, kernel_size, True, layer_idx=i * num_blocks + j))
# Final projection to spectrogram dimensions (magnitude + phase)
layers.append(nn.Conv1d(hidden_size, n_fft + 2, 1, bias=False))
self.layers = nn.ModuleList(layers)
def forward(self, x: Tensor, cache=None, flush: bool = False, constrain_value_range: bool = True) -> Tensor:
# Transpose input from [batch, time, channels] to [batch, channels, time]
x = x.transpose(-1, -2)
# Pass through the network
for layer in self.layers:
if isinstance(layer, ConvNeXt1d):
x = layer(x, cache=cache, flush=flush)
else:
x = layer(x)
# Convert network output to a complex spectrogram and then to a waveform
with torch.autocast(device_type=x.device.type, enabled=False):
max_mag = 100.0
# Split output into magnitude and phase components
mag, ph = x.float().chunk(2, dim=1)
# Safeguard to prevent excessively large magnitudes
mag = max_mag * torch.exp(-F.softplus(-mag + math.log(max_mag)))
# Reconstruct the complex spectrogram from magnitude and phase
# The DC and Nyquist components are real, so their phase is applied via cosine.
mag_dc, mag_mid, mag_nyquist = mag.split([1, mag.size(1) - 2, 1], dim=1)
ph_dc, ph_mid, ph_nyquist = torch.cos(ph).split([1, ph.size(1) - 2, 1], dim=1)
ph_imag = torch.sin(ph[:, 1:-1, :])
spec_real = mag_mid * ph_mid
spec_imag = mag_mid * ph_imag
spec = torch.cat([mag_dc * ph_dc, spec_real + 1j * spec_imag, mag_nyquist * ph_nyquist], 1)
# Handle caching for autoregressive generation of the spectrogram
if cache is not None:
half_spec_padding = math.ceil(((self.n_fft - self.hop_length) // 2) / self.hop_length)
spec = cache.update(spec, self.spec_cache_idx, padding=half_spec_padding * 2, flush=flush)
if flush:
spec = F.pad(spec, [0, half_spec_padding])
# Convert spectrogram to waveform
x = spec_to_wav(
spec, self.n_fft, self.hop_length, self.n_fft, constrain_value_range=constrain_value_range
).unsqueeze(1)
if cache is not None:
# Trim the output to remove the padded region from the start
half_wav_padding = half_spec_padding * self.hop_length
x = x[:, :, half_wav_padding:-half_wav_padding]
return x
class RVQVAEModel(nn.Module):
"""
Residual Vector-Quantized Variational Autoencoder (RVQ-VAE) model.
This model learns a discrete representation of audio by encoding a waveform
into a latent space and then quantizing the latents into discrete codes.
It consists of an encoder, a quantizer, and a decoder.
Args:
config (DictConfig | dict[str, Any]): A configuration object with model hyperparameters.
"""
config_class: type[DictConfig] = DictConfig
def __init__(self, config: DictConfig | dict[str, Any]):
super().__init__()
self.config = config
self.encoder = Wav2Latent(
latent_size=self.config.latent_size,
n_fft=self.config.n_fft,
hop_length=self.config.hop_length,
base_hidden_size=self.config.base_hidden_size,
channel_mult=self.config.channel_mult,
rates=self.config.rates,
num_blocks=self.config.num_blocks,
kernel_size=self.config.kernel_size,
groups=self.config.groups,
)
# Layers for quantization
self.prvq = PreTrainedProbabilisticVQ(
channels=self.config.latent_size,
num_mixtures=self.config.codebook_size,
depth=self.config.num_quantizers,
)
self.decoder = Latent2Wav(
latent_size=self.config.latent_size,
n_fft=self.config.n_fft,
hop_length=self.config.hop_length,
base_hidden_size=self.config.base_hidden_size,
channel_mult=tuple(reversed(self.config.channel_mult)),
rates=tuple(reversed(self.config.rates)),
num_blocks=self.config.num_blocks,
kernel_size=self.config.kernel_size,
groups=self.config.groups,
)
for p in self.parameters():
p.requires_grad = False
def ae_encode(self, x: Tensor, cache: CausalConv1dCache | None = None, flush: bool = False) -> Tensor:
"""
Runs the encoder part of the autoencoder.
Args:
x (Tensor): Input waveform. Shape: `[batch, 1, time]`.
cache (CausalConv1dCache | None): Not implemented for the encoder.
flush (bool): Not implemented for the encoder.
Returns:
Tensor: The continuous latent representation. Shape: `[batch, time', channels]`.
"""
assert x.size(1) == 1 and x.dim() == 3, "Input must be a batch of mono audio."
assert x.size(2) % self.config.wav_to_token_ratio == 0, (
f"Input audio length ({x.size(2)}) must be divisible by the model's "
f"wav_to_token_ratio ({self.config.wav_to_token_ratio}). "
f"Please pad the input to a compatible length."
)
if cache is not None:
raise NotImplementedError("Caching is not supported for the encoder.")
return self.encoder(x, cache=cache, flush=flush)
def ae_decode(
self,
x: Tensor,
constrain_value_range: bool = True,
cache: CausalConv1dCache | None = None,
flush: bool = False,
) -> Tensor:
"""
Runs the decoder part of the autoencoder.
Args:
x (Tensor): The (de-quantized) latent representation. Shape: `[batch, time', channels]`.
constrain_value_range (bool): If True, constrains the output of the ISTFT.
cache (CausalConv1dCache | None): Cache for autoregressive generation.
flush (bool): If True, flushes the cache.
Returns:
Tensor: The reconstructed waveform. Shape: `[batch, 1, time]`.
"""
return self.decoder(x, constrain_value_range=constrain_value_range, cache=cache, flush=flush)
def encode(self, x: Tensor, x_len: Tensor) -> tuple[Tensor, Tensor]:
"""
Encodes a waveform into discrete codes.
Args:
x (Tensor): Input waveform. Shape: `[batch, 1, time]`.
x_len (Tensor): The original lengths of the waveforms in the batch.
Returns:
tuple[Tensor, Tensor]: A tuple containing:
- The discrete codes. Shape: `[batch, time', n_quantizers]`.
- The lengths of the code sequences.
"""
with disable_tf32():
z_e = self.ae_encode(x)
code_len = x_len // self.config.wav_to_token_ratio
return self.quantize(z_e), code_len
def decode(
self,
code: Tensor,
code_len: Tensor | None = None,
constrain_value_range: bool = True,
cache: CausalConv1dCache | None = None,
flush: bool = False,
) -> tuple[Tensor, Tensor | None]:
"""
Decodes discrete codes back into a waveform.
Args:
code (Tensor): The discrete codes. Shape: `[batch, time', n_quantizers]`.
code_len (Tensor | None): The lengths of the code sequences.
constrain_value_range (bool): If True, constrains the output of the ISTFT.
cache (CausalConv1dCache | None): Cache for autoregressive generation.
flush (bool): If True, flushes the cache.
Returns:
tuple[Tensor, Tensor | None]: A tuple containing:
- The reconstructed waveform. Shape: `[batch, 1, time]`.
- The lengths of the reconstructed waveforms.
"""
with disable_tf32():
z_q = self.dequantize(code)
x_hat = self.ae_decode(z_q, constrain_value_range=constrain_value_range, cache=cache, flush=flush)
wav_len = code_len * self.config.wav_to_token_ratio if code_len is not None else None
return x_hat, wav_len
def quantize(self, z: Tensor) -> Tensor:
"""
Quantizes a continuous latent tensor into discrete codes.
Args:
z (Tensor): The continuous latent tensor from the encoder.
Shape: `[batch, time, channels]`.
Returns:
Tensor: The quantized codes. Shape: `[batch, time, n_quantizers]`.
"""
with disable_tf32():
ids_sel = self.prvq.encode(z, return_z_q=False)
return torch.stack(ids_sel, -1)
def dequantize(self, code: Tensor) -> Tensor:
"""
De-quantizes discrete codes back into a continuous latent tensor.
Args:
code (Tensor): The quantized codes. Shape: `[batch, time, n_quantizers]`.
Returns:
Tensor: The de-quantized continuous latent tensor.
Shape: `[batch, time, latent_size]`.
"""
ids_sel = [x.squeeze(-1) for x in torch.split(code, 1, -1)]
return self.prvq.decode(ids_sel)
def forward(self, x: Tensor, constrain_value_range: bool = False) -> Tensor:
"""
Performs a full autoencoding pass: encode, quantize, dequantize, and decode.
Args:
x (Tensor): The input waveform. Shape: `[batch, 1, time]`.
constrain_value_range (bool): If True, constrains the output of the ISTFT.
Returns:
Tensor: The reconstructed waveform. Shape: `[batch, 1, time]`.
"""
with torch.no_grad():
z_e = self.ae_encode(x)
code = self.quantize(z_e)
z_d = self.dequantize(code)
x_hat = self.ae_decode(z_d, constrain_value_range=constrain_value_range)
return x_hat
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/fusion.py
+403
View File
@@ -0,0 +1,403 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/speechlm2/parts/fusion.py; training-only code removed.
"""
Fusion modules for combining multi-modal embeddings in duplex speech models.
Supported fusion methods:
- "add" or None: Simple additive fusion with per-modality weights
- "concat": Concatenate (after LayerNorm) and project back to hidden_dim
- "gated_simple": Per-timestep, per-dimension gates with softmax normalization
- "gated_gmu": Gated Multimodal Unit (Arevalo et al., 2017) with per-modality sigmoid gates
For gated methods, the channel embeddings, tie_and_roll, and per-modality weights
are bypassed - the gating mechanism learns the optimal combination.
"""
import torch
import torch.nn as nn
from typing import Optional, Literal, Union
from .compat import logging
FuseMethod = Literal["add", "concat", "gated_simple", "gated_gmu"]
class FusionModule(nn.Module):
"""Base class for fusion modules."""
def __init__(self, hidden_dim: int):
super().__init__()
self.hidden_dim = hidden_dim
def forward(
self,
agent_text_embeds: torch.Tensor,
user_audio_embeds: torch.Tensor,
user_text_embeds: torch.Tensor | None = None,
function_embeds: torch.Tensor | None = None,
) -> torch.Tensor:
"""
Fuse multi-modal embeddings.
Args:
agent_text_embeds: Agent text embeddings (B, T, D)
user_audio_embeds: User audio embeddings (B, T, D)
user_text_embeds: User text/ASR embeddings (B, T, D) or None
function_embeds: Function calling channel embeddings (B, T, D) or None
Returns:
Fused embeddings (B, T, D)
"""
raise NotImplementedError
class AddFusion(FusionModule):
"""
Simple additive fusion with optional per-modality weights.
This is the default behavior matching the original implementation.
"""
def __init__(
self,
hidden_dim: int,
agent_text_weight: float = 1.0,
user_audio_weight: float = 1.0,
user_text_weight: float = 1.0,
function_weight: float = 1.0,
):
super().__init__(hidden_dim)
self.agent_text_weight = agent_text_weight
self.user_audio_weight = user_audio_weight
self.user_text_weight = user_text_weight
self.function_weight = function_weight
def forward(
self,
agent_text_embeds: torch.Tensor,
user_audio_embeds: torch.Tensor,
user_text_embeds: torch.Tensor | None = None,
function_embeds: torch.Tensor | None = None,
) -> torch.Tensor:
output = agent_text_embeds * self.agent_text_weight
output = output + user_audio_embeds * self.user_audio_weight
if user_text_embeds is not None:
output = output + user_text_embeds * self.user_text_weight
if function_embeds is not None:
output = output + function_embeds * self.function_weight
return output
class ConcatFusion(FusionModule):
"""
Concatenate embeddings (after LayerNorm) and project back to hidden_dim.
Supports optional user_text_embeds and function_embeds; zeros are used when
not provided. Uses 4 modalities (agent, audio, user_text, function) so that
function channel can be optionally present.
The linear layer learns to handle this case appropriately.
"""
def __init__(
self,
hidden_dim: int,
agent_text_weight: float = 1.0,
user_audio_weight: float = 1.0,
user_text_weight: float = 1.0,
function_weight: float = 1.0,
num_modalities: int = 4,
):
super().__init__(hidden_dim)
self.agent_text_weight = agent_text_weight
self.user_audio_weight = user_audio_weight
self.user_text_weight = user_text_weight
self.function_weight = function_weight
self.num_modalities = num_modalities
# LayerNorm for each modality
self.ln_agent = nn.LayerNorm(hidden_dim)
self.ln_audio = nn.LayerNorm(hidden_dim)
self.ln_text = nn.LayerNorm(hidden_dim)
self.ln_function = nn.LayerNorm(hidden_dim)
# Projection layer (4 modalities; pad with zeros when optional channels missing)
self.proj = nn.Linear(hidden_dim * num_modalities, hidden_dim, bias=False)
def forward(
self,
agent_text_embeds: torch.Tensor,
user_audio_embeds: torch.Tensor,
user_text_embeds: torch.Tensor | None = None,
function_embeds: torch.Tensor | None = None,
) -> torch.Tensor:
# Apply layernorm and weights
agent_normed = self.ln_agent(agent_text_embeds) * self.agent_text_weight
audio_normed = self.ln_audio(user_audio_embeds) * self.user_audio_weight
if user_text_embeds is not None:
text_normed = self.ln_text(user_text_embeds) * self.user_text_weight
else:
# Zero-pad when ASR embeddings are not available
text_normed = torch.zeros_like(agent_text_embeds)
if function_embeds is not None:
function_normed = self.ln_function(function_embeds) * self.function_weight
else:
function_normed = torch.zeros_like(agent_text_embeds)
# Concatenate along hidden dim and project
concat = torch.cat([agent_normed, audio_normed, text_normed, function_normed], dim=-1)
return self.proj(concat)
class GatedFusionSimple(FusionModule):
"""
Learns a softmax-normalized gate over 4 streams at each timestep (agent, audio, user_text, function).
Gate input: LayerNorm(each stream) → concatenate → linear → softmax
Output: weighted sum of LayerNorm'd embeddings where weights are input-dependent.
Optional channels (user_text, function) use zeros when not provided; their gate is masked to 0.
Note: This is PER-TIMESTEP gating (one scalar weight per modality per timestep).
For PER-DIMENSION gating, use GatedFusionGMU.
"""
def __init__(self, hidden_dim: int, num_modalities: int = 4):
super().__init__(hidden_dim)
self.num_modalities = num_modalities
# LayerNorm for each stream (normalizes distributions before gate computation)
self.ln_agent = nn.LayerNorm(hidden_dim)
self.ln_audio = nn.LayerNorm(hidden_dim)
self.ln_text = nn.LayerNorm(hidden_dim)
self.ln_function = nn.LayerNorm(hidden_dim)
# Gate network: outputs one gate value per modality per timestep
# Input: concatenated normalized modalities (B, T, D*4)
# Output: gates (B, T, 4) - one scalar per modality
self.gate_net = nn.Linear(hidden_dim * num_modalities, num_modalities, bias=True)
# Initialize for uniform weighting
self._init_uniform()
def _init_uniform(self):
"""Initialize gate network for uniform initial weighting."""
nn.init.zeros_(self.gate_net.weight)
nn.init.zeros_(self.gate_net.bias)
def forward(
self,
agent_text_embeds: torch.Tensor,
user_audio_embeds: torch.Tensor,
user_text_embeds: torch.Tensor | None = None,
function_embeds: torch.Tensor | None = None,
) -> torch.Tensor:
B, T, D = agent_text_embeds.shape
has_user_text = user_text_embeds is not None
has_function = function_embeds is not None
text_input = user_text_embeds if has_user_text else torch.zeros_like(agent_text_embeds)
function_input = function_embeds if has_function else torch.zeros_like(agent_text_embeds)
# Apply LayerNorm to each stream before gate computation
agent_normed = self.ln_agent(agent_text_embeds)
audio_normed = self.ln_audio(user_audio_embeds)
text_normed = self.ln_text(text_input)
function_normed = self.ln_function(function_input)
# Concatenate normalized inputs for gate computation
concat = torch.cat([agent_normed, audio_normed, text_normed, function_normed], dim=-1)
# Compute gates: (B, T, D*4) -> (B, T, 4)
gates = self.gate_net(concat) # (B, T, num_modalities)
# Mask out optional channels when not present (before softmax)
if not has_user_text:
gates = gates.clone()
gates[:, :, 2] = float("-inf")
if not has_function:
gates = gates.clone()
gates[:, :, 3] = float("-inf")
# Softmax over modalities for normalization
gates = torch.softmax(gates, dim=-1) # (B, T, 4)
# Stack normalized modalities for weighted sum: (B, T, 4, D)
stacked = torch.stack([agent_normed, audio_normed, text_normed, function_normed], dim=2)
# Apply gates (broadcast) and sum over modalities
fused = (stacked * gates.unsqueeze(-1)).sum(dim=2) # (B, T, D)
return fused
class GatedFusionGMU(FusionModule):
"""
Gated Multimodal Unit inspired by Arevalo et al. (2017).
PER-DIMENSION gating: each dimension can attend to different modalities.
Uses softmax over modalities for numerically stable normalization.
Supports 4 modalities (agent, audio, user_text, function); optional channels use zeros and are masked.
All streams are LayerNorm'd before gate computation and fusion.
Gate output shape: (B, T, 4, D) with softmax over the modality dimension.
"""
def __init__(self, hidden_dim: int, num_modalities: int = 4):
super().__init__(hidden_dim)
self.num_modalities = num_modalities
# LayerNorm for each stream (normalizes distributions before gate computation)
self.ln_agent = nn.LayerNorm(hidden_dim)
self.ln_audio = nn.LayerNorm(hidden_dim)
self.ln_text = nn.LayerNorm(hidden_dim)
self.ln_function = nn.LayerNorm(hidden_dim)
# Per-modality gate projections (output logits, not activations)
self.gate_agent = nn.Linear(hidden_dim * num_modalities, hidden_dim)
self.gate_audio = nn.Linear(hidden_dim * num_modalities, hidden_dim)
self.gate_text = nn.Linear(hidden_dim * num_modalities, hidden_dim)
self.gate_function = nn.Linear(hidden_dim * num_modalities, hidden_dim)
# Initialize for uniform weighting (all logits = 0 → equal softmax weights)
self._init_uniform()
def _init_uniform(self):
"""Initialize gate networks for uniform initial weighting."""
for gate in [self.gate_agent, self.gate_audio, self.gate_text, self.gate_function]:
nn.init.zeros_(gate.weight)
nn.init.zeros_(gate.bias)
def forward(
self,
agent_text_embeds: torch.Tensor,
user_audio_embeds: torch.Tensor,
user_text_embeds: torch.Tensor | None = None,
function_embeds: torch.Tensor | None = None,
) -> torch.Tensor:
B, T, D = agent_text_embeds.shape
has_user_text = user_text_embeds is not None
has_function = function_embeds is not None
text_input = user_text_embeds if has_user_text else torch.zeros_like(agent_text_embeds)
function_input = function_embeds if has_function else torch.zeros_like(agent_text_embeds)
# Apply LayerNorm to each stream before gate computation
agent_normed = self.ln_agent(agent_text_embeds)
audio_normed = self.ln_audio(user_audio_embeds)
text_normed = self.ln_text(text_input)
function_normed = self.ln_function(function_input)
# Concatenate normalized inputs for gate computation
concat = torch.cat([agent_normed, audio_normed, text_normed, function_normed], dim=-1)
# Compute per-modality gate logits: each is (B, T, D)
logits_agent = self.gate_agent(concat)
logits_audio = self.gate_audio(concat)
logits_text = self.gate_text(concat)
logits_function = self.gate_function(concat)
# Stack logits: (B, T, 4, D)
stacked_logits = torch.stack([logits_agent, logits_audio, logits_text, logits_function], dim=2)
# Mask out optional channels when not present (before softmax)
if not has_user_text:
stacked_logits = stacked_logits.clone()
stacked_logits[:, :, 2, :] = float("-inf")
if not has_function:
stacked_logits = stacked_logits.clone()
stacked_logits[:, :, 3, :] = float("-inf")
# Softmax over modalities (dim=2) for stable normalization
gates = torch.softmax(stacked_logits, dim=2) # (B, T, 4, D)
# Stack normalized embeddings: (B, T, 4, D)
stacked_embeds = torch.stack([agent_normed, audio_normed, text_normed, function_normed], dim=2)
# Fuse with element-wise gating
fused = (gates * stacked_embeds).sum(dim=2) # (B, T, D)
return fused
def create_fusion_module(
fuse_method: Union[str, FuseMethod] | None,
hidden_dim: int,
agent_text_weight: float = 1.0,
user_audio_weight: float = 1.0,
user_text_weight: float = 1.0,
function_weight: float = 1.0,
) -> FusionModule:
"""
Factory function to create the appropriate fusion module.
Args:
fuse_method: One of None, "add", "concat", "gated_simple", "gated_gmu"
hidden_dim: Hidden dimension of the embeddings
agent_text_weight: Weight for agent text (used by add/concat methods only)
user_audio_weight: Weight for user audio (used by add/concat methods only)
user_text_weight: Weight for user text (used by add/concat methods only)
function_weight: Weight for function channel (used by add/concat only; optional channel)
Returns:
Instantiated fusion module
Notes:
- All fusion modules accept optional function_embeds; when None, that channel
contributes nothing (zeros + mask for gated methods).
- For gated methods (gated_simple, gated_gmu), the weight parameters are
ignored as the gating mechanism learns the optimal combination.
- When using gated methods, you should also bypass channel embeddings,
tie_and_roll, and any manual embedding weight parameters in the model.
"""
if fuse_method is None or fuse_method == "add":
logging.info(
f"Using AddFusion with weights: agent={agent_text_weight}, "
f"audio={user_audio_weight}, text={user_text_weight}, function={function_weight}"
)
return AddFusion(
hidden_dim=hidden_dim,
agent_text_weight=agent_text_weight,
user_audio_weight=user_audio_weight,
user_text_weight=user_text_weight,
function_weight=function_weight,
)
elif fuse_method == "concat":
logging.info(
f"Using ConcatFusion with weights: agent={agent_text_weight}, "
f"audio={user_audio_weight}, text={user_text_weight}, function={function_weight}"
)
return ConcatFusion(
hidden_dim=hidden_dim,
agent_text_weight=agent_text_weight,
user_audio_weight=user_audio_weight,
user_text_weight=user_text_weight,
function_weight=function_weight,
)
elif fuse_method == "gated_simple":
logging.info(
f"Using GatedFusionSimple with hidden_dim={hidden_dim} (learned gating, optional function channel)"
)
return GatedFusionSimple(hidden_dim=hidden_dim)
elif fuse_method == "gated_gmu":
logging.info(f"Using GatedFusionGMU with hidden_dim={hidden_dim} (learned gating, optional function channel)")
return GatedFusionGMU(hidden_dim=hidden_dim)
else:
raise ValueError(
f"Unknown fuse_method: '{fuse_method}'. Expected one of: None, 'add', 'concat', 'gated_simple', 'gated_gmu'"
)
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/perception.py
+135
View File
@@ -0,0 +1,135 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat) nemo/collections/speechlm2/modules/perception.py; training-only code removed.
import torch
from omegaconf import DictConfig
from torch import nn
from .compat import (
Exportable,
NeuralModule,
typecheck,
)
class AudioPerceptionModule(NeuralModule, Exportable):
"""Audio perception module that consists of audio encoder(s) and modality adapter."""
def __init__(self, cfg: DictConfig):
super().__init__()
# Initialize components
self.cfg = cfg
self.preprocessor = self.from_config_dict(cfg.preprocessor)
self.encoder = self.from_config_dict(cfg.encoder)
if "spec_augment" in cfg and cfg.spec_augment is not None:
self.spec_augmentation = self.from_config_dict(cfg.spec_augment)
else:
self.spec_augmentation = None
self.modality_adapter = self.from_config_dict(cfg.modality_adapter)
if "output_dim" not in cfg.modality_adapter and "d_model" in cfg.modality_adapter: # e.g., conformer encoder
self.proj = nn.Linear(cfg.modality_adapter.d_model, cfg.output_dim)
else:
self.proj = nn.Identity()
self.modality_adapter_quantizer_levels = cfg.get("modality_adapter_quantizer_levels", None)
if self.modality_adapter_quantizer_levels:
# FiniteScalarQuantizer (nemo.collections.tts.modules.audio_codec_modules) was not
# vendored: `modality_adapter_quantizer_levels` is unset in the
# NVIDIA-NemotronLabs-VoiceChat-11B checkpoint config.json.
raise RuntimeError(
"modality_adapter_quantizer_levels is set in the perception config, but "
"FiniteScalarQuantizer was not vendored (it is unused by the "
"NVIDIA-NemotronLabs-VoiceChat-11B checkpoint). Vendor "
"nemo.collections.tts.modules.audio_codec_modules.FiniteScalarQuantizer to enable it."
)
def maybe_preprocess_audio(
self,
input_signal=None,
input_signal_length=None,
processed_signal=None,
processed_signal_length=None,
):
has_input_signal = input_signal is not None and input_signal_length is not None
has_processed_signal = processed_signal is not None and processed_signal_length is not None
if (has_input_signal ^ has_processed_signal) is False:
raise ValueError(
f"{self.__class__} Arguments ``input_signal`` and ``input_signal_length`` are mutually exclusive "
" with ``processed_signal`` and ``processed_signal_len`` arguments."
)
if not has_processed_signal:
processed_signal, processed_signal_length = self.preprocessor(
input_signal=input_signal,
length=input_signal_length,
)
return processed_signal, processed_signal_length
# disable type checks to avoid type-check errors when using Conformer as modality adapter
@typecheck.disable_checks()
def forward(
self,
input_signal=None,
input_signal_length=None,
processed_signal=None,
processed_signal_length=None,
return_encoder_emb=False,
sample_id=None,
):
processed_signal, processed_signal_length = self.maybe_preprocess_audio(
input_signal, input_signal_length, processed_signal, processed_signal_length
)
# Spec augment is not applied during evaluation/testing
if self.spec_augmentation is not None and self.training:
processed_signal = self.spec_augmentation(input_spec=processed_signal, length=processed_signal_length)
encoder_emb, encoded_len = self.encoder(audio_signal=processed_signal, length=processed_signal_length)
if self.modality_adapter_quantizer_levels is not None:
encoded = self.modality_adapter_quantizer_bottleneck(encoder_emb.transpose(1, 2))
encoded, _ = self.modality_adapter_vector_quantizer(inputs=encoded.transpose(1, 2), input_len=None)
encoded = self.modality_adapter_quantizer_projection(encoded.transpose(1, 2)).transpose(1, 2)
else:
encoded = encoder_emb
encoded, encoded_len = self.modality_adapter(audio_signal=encoded, length=encoded_len)
# encoded_orig = encoded.clone()
# b, c, t -> b, t, c
encoded = self.proj(encoded.transpose(1, 2))
if return_encoder_emb:
return encoded, encoded_len, encoder_emb.transpose(1, 2)
else:
return encoded, encoded_len
class IdentityConnector(NeuralModule, Exportable):
"""User to pass encoder's representations as-is to the LLM."""
def __init__(
self,
*args,
**kwargs,
):
super().__init__()
def forward(self, audio_signal, length=None, *args, **kwargs):
return audio_signal, length
vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/tokenizer.py
+501
View File
@@ -0,0 +1,501 @@
# SPDX-FileCopyrightText: Copyright (c) 2020 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Vendored from NVIDIA-NeMo/Speech (branch nemotron-labs-voicechat)
# nemo/collections/common/tokenizers/tokenizer_spec.py and
# nemo/collections/common/tokenizers/huggingface/auto_tokenizer.py; training-only code removed.
# Only the imports were adjusted (nemo.utils logging -> stdlib shim); the class
# bodies are exact copies.
from abc import ABC, abstractmethod
from collections import OrderedDict
from typing import List, Optional
from transformers import AutoTokenizer as AUTOTOKENIZER
from .compat import logging
__all__ = ["TokenizerSpec", "AutoTokenizer"]
class TokenizerSpec(ABC):
"""
Inherit this class to implement a new tokenizer.
"""
@abstractmethod
def text_to_tokens(self, text):
"""Converts text into a list of tokens."""
pass
@abstractmethod
def tokens_to_ids(self, tokens):
"""Converts a list of tokens to their corresponding IDs."""
pass
@abstractmethod
def ids_to_tokens(self, ids):
"""Converts a list of token IDs back to tokens."""
pass
@abstractmethod
def text_to_ids(self, text):
"""Converts text directly to token IDs."""
pass
def add_special_tokens(self, special_tokens: List[str]):
"""Adds special tokens (eos, pad, cls...) to vocab."""
raise NotImplementedError("To be implemented")
def apply_chat_template(self, *args, **kwargs):
"""Applies chat template and tokenizes results"""
raise NotImplementedError("To be implemented")
@property
def name(self):
"""name of the class"""
return type(self).__name__
@property
def cls(self):
"""Property alias to match MegatronTokenizer; returns cls_id if available."""
if hasattr(self, "cls_id"):
return self.cls_id
raise AttributeError(f"{type(self).__name__} has no attribute 'cls' or 'cls_id'")
@property
def sep(self):
"""Property alias to match MegatronTokenizer; returns sep_id if available."""
if hasattr(self, "sep_id"):
return self.sep_id
raise AttributeError(f"{type(self).__name__} has no attribute 'sep' or 'sep_id'")
@property
def pad(self):
"""Property alias to match MegatronTokenizer; returns pad_id if available."""
if hasattr(self, "pad_id"):
return self.pad_id
raise AttributeError(f"{type(self).__name__} has no attribute 'pad' or 'pad_id'")
@property
def eod(self):
"""Property alias to match MegatronTokenizer; returns eod_id if available."""
if hasattr(self, "eod_id"):
return self.eod_id
if hasattr(self, "eos_id"):
# Default to end-of-sentence id if end-of-document is not defined.
return self.eos_id
raise AttributeError(f"{type(self).__name__} has no attribute 'eod', 'eod_id', 'eos', or 'eos_id'")
@property
def bos(self):
"""Property alias to match MegatronTokenizer; returns bos_id if available."""
if hasattr(self, "bos_id"):
return self.bos_id
raise AttributeError(f"{type(self).__name__} has no attribute 'bos' or 'bos_id'")
@property
def eos(self):
"""Property alias to match MegatronTokenizer; returns eos_id if available."""
if hasattr(self, "eos_id"):
return self.eos_id
raise AttributeError(f"{type(self).__name__} has no attribute 'eos' or 'eos_id'")
@property
def mask(self):
"""Property alias to match MegatronTokenizer; returns mask_id if available."""
if hasattr(self, "mask_id"):
return self.mask_id
raise AttributeError(f"{type(self).__name__} has no attribute 'mask' or 'mask_id'")
class AutoTokenizer(TokenizerSpec):
"""
Wrapper of HuggingFace AutoTokenizer
https://huggingface.co/docs/transformers/main/en/model_doc/auto#transformers.AutoTokenizer.
"""
def __init__(
self,
pretrained_model_name: str,
vocab_file: str | None = None,
merges_file: str | None = None,
mask_token: str | None = None,
bos_token: str | None = None,
eos_token: str | None = None,
pad_token: str | None = None,
sep_token: str | None = None,
cls_token: str | None = None,
unk_token: str | None = None,
additional_special_tokens: List | None = [],
use_fast: bool | None = True,
trust_remote_code: bool | None = False,
include_special_tokens: bool = False,
):
"""
Args:
pretrained_model_name: corresponds to HuggingFace-AutoTokenizer's 'pretrained_model_name_or_path' input
argument. For more details please refer to the documentation of the `from_pretrained` method here:
https://huggingface.co/docs/transformers/main/en/model_doc/auto#transformers.AutoTokenizer.
The list of all supported models can be found here: https://huggingface.co/models
vocab_file: path to file with vocabulary which consists
of characters separated by newlines.
mask_token: mask token
bos_token: the beginning of sequence token
eos_token: the end of sequence token. Usually equal to sep_token
pad_token: token to use for padding
sep_token: token used for separating sequences
cls_token: class token. Usually equal to bos_token
unk_token: token to use for unknown tokens
additional_special_tokens: list of other tokens beside standard special tokens (bos, eos, pad, etc.). For
example, sentinel tokens for T5 (<extra_id_0>, <extra_id_1>, etc.)
use_fast: whether to use fast HuggingFace tokenizer
include_special_tokens: when True, converting text to ids will include special tokens / prompt tokens (if
any), yielding self.tokenizer(text).input_ids
"""
try:
self._initialize_tokenizer(pretrained_model_name, vocab_file, merges_file, use_fast, trust_remote_code)
assert self.tokenizer, "tokenizer not initialized"
except Exception:
try:
self._initialize_tokenizer(
pretrained_model_name, vocab_file, merges_file, not use_fast, trust_remote_code
)
assert self.tokenizer, "tokenizer not initialized"
except Exception as e:
raise ValueError(
f"Unable to instantiate HuggingFace AUTOTOKENIZER for {pretrained_model_name}. Exception: {e}"
)
self.include_special_tokens = include_special_tokens
self.original_vocab_size = len(self.tokenizer)
special_tokens_dict = {}
# # setting special tokens, by default the default model's special tokens will be preserved
# # unless passes new values to the special tokens
if unk_token is not None:
special_tokens_dict["unk_token"] = unk_token
if mask_token is not None:
special_tokens_dict["mask_token"] = mask_token
if pad_token is not None:
special_tokens_dict["pad_token"] = pad_token
# if the model does not have eos_token but has sep_token,
# set eos_token = sep_token, and vice versa
if sep_token is not None:
special_tokens_dict["sep_token"] = sep_token
elif self.tokenizer.sep_token is None and self.tokenizer.eos_token:
special_tokens_dict["sep_token"] = self.tokenizer.eos_token
if eos_token is not None:
special_tokens_dict["eos_token"] = eos_token
elif self.tokenizer.eos_token is None and self.tokenizer.sep_token:
special_tokens_dict["eos_token"] = self.tokenizer.sep_token
# if the model does not have bos_token but has cls_token,
# set bos_token = cls_token, and vice versa
if bos_token is not None:
special_tokens_dict["bos_token"] = bos_token
elif self.tokenizer.bos_token is None and self.tokenizer.cls_token:
special_tokens_dict["bos_token"] = self.tokenizer.cls_token
if cls_token is not None:
special_tokens_dict["cls_token"] = cls_token
elif self.tokenizer.cls_token is None and self.tokenizer.bos_token:
special_tokens_dict["cls_token"] = self.tokenizer.bos_token
# add additional special tokens (not standard special tokens such as bos, eod, sep)
if additional_special_tokens is not None:
special_tokens_dict["additional_special_tokens"] = additional_special_tokens
new_tokens_in_vocab = []
for token in [mask_token, bos_token, eos_token, pad_token, sep_token, cls_token, unk_token]:
if token is not None and token not in self.tokenizer.get_vocab():
new_tokens_in_vocab.append(token)
for token in additional_special_tokens:
if token is not None and token not in self.tokenizer.get_vocab():
new_tokens_in_vocab.append(token)
if len(new_tokens_in_vocab) > 0:
"""
Special tokens that were not previously included in the tokenizer's vocabulary file will be added to
the vocabulary and, as a result, the model should be resized, for example:
# define your model
pretrained_model_name = 'roberta-base'
model = nemo_nlp.modules.get_lm_model(pretrained_model_name=pretrained_model_name)
# define pretrained tokenizer
tokenizer_default = nemo_nlp.modules.get_tokenizer(tokenizer_name=pretrained_model_name)
special_tokens = {'bos_token': '<BOS>',
'cls_token': '<CSL>',
'additional_special_tokens': ['<MY_NER_TOKEN>', '<ANOTHER_TOKEN>']}
tokenizer_default.add_special_tokens(special_tokens_dict=special_tokens)
# resize your model so that the embeddings for newly added tokens are updated during training/finetuning
model.resize_token_embeddings(tokenizer_default.vocab_size)
See NLP_Tokenizers.ipynb for more details.
"""
logging.warning(
f"{new_tokens_in_vocab} \n will be added to the vocabulary.\n"
f"Please resize your model accordingly, "
f"see NLP_Tokenizers.ipynb for more details."
)
self.add_special_tokens(special_tokens_dict)
self.space_sensitive = self.text_to_tokens("x y") != self.text_to_tokens("x") + self.text_to_tokens("y")
self._inv_vocab_dict = {}
def _initialize_tokenizer(
self,
pretrained_model_name: str,
vocab_file: str | None = None,
merges_file: str | None = None,
use_fast: bool | None = False,
trust_remote_code: bool | None = False,
):
# this logic deals with different huggingface tokenizers having different positional args
if vocab_file is None:
self.tokenizer = AUTOTOKENIZER.from_pretrained(
pretrained_model_name_or_path=pretrained_model_name,
use_fast=use_fast,
trust_remote_code=trust_remote_code,
)
elif merges_file is None:
self.tokenizer = AUTOTOKENIZER.from_pretrained(
pretrained_model_name_or_path=pretrained_model_name,
vocab_file=vocab_file,
use_fast=use_fast,
trust_remote_code=trust_remote_code,
)
else:
self.tokenizer = AUTOTOKENIZER.from_pretrained(
pretrained_model_name_or_path=pretrained_model_name,
vocab_file=vocab_file,
merges_file=merges_file,
use_fast=use_fast,
trust_remote_code=trust_remote_code,
)
@property
def vocab_size(self):
"""
Returns the size of the tokenizer's vocabulary.
Returns:
int: The number of tokens in the vocabulary.
"""
return len(self.tokenizer)
def add_special_tokens(self, special_tokens_dict: dict) -> int:
"""
Adds a dictionary of special tokens (eos, pad, cls...). If special tokens are NOT in the vocabulary, they are
added to it (indexed starting from the last index of the current vocabulary).
Args:
special_tokens_dict: dict of string. Keys should be in the list of predefined special attributes:
[``bos_token``, ``eos_token``, ``unk_token``, ``sep_token``, ``pad_token``, ``cls_token``,
``mask_token``, ``additional_special_tokens``].
Tokens are only added if they are not already in the vocabulary.
Returns:
Number of tokens added to the vocabulary.
"""
num_tokens_added = self.tokenizer.add_special_tokens(special_tokens_dict)
if num_tokens_added > 0:
logging.info(f"{num_tokens_added} special tokens added, resize your model accordingly.")
for k in self.tokenizer.SPECIAL_TOKENS_ATTRIBUTES:
setattr(self, k, getattr(self.tokenizer, k, None))
return num_tokens_added
def text_to_tokens(self, text):
"""
Converts text into a list of tokens.
Args:
text (str): Input text to be tokenized.
Returns:
List[str]: List of tokens.
"""
tokens = self.tokenizer.tokenize(text)
return tokens
def tokens_to_ids(self, tokens):
"""
Converts a list of tokens to their corresponding IDs.
Args:
tokens (List[str]): List of tokens to convert.
Returns:
List[int]: List of token IDs.
"""
ids = self.tokenizer.convert_tokens_to_ids(tokens)
return ids
def ids_to_tokens(self, ids):
"""
Converts a list of token IDs back to tokens.
Args:
ids (List[int]): List of token IDs to convert.
Returns:
List[str]: List of tokens.
"""
tokens = self.tokenizer.convert_ids_to_tokens(ids)
return tokens
def text_to_ids(self, text):
"""
Converts text directly to token IDs.
Args:
text (str): Input text to be converted to IDs.
Returns:
List[int]: List of token IDs. If include_special_tokens is True, will include special tokens from the
tokenizer's configuration.
"""
if self.include_special_tokens:
return self.tokenizer(text).input_ids
tokens = self.text_to_tokens(text)
ids = self.tokens_to_ids(tokens)
return ids
def apply_chat_template(self, *args, **kwargs):
"""Applies chat template and tokenizes results"""
return self.tokenizer.apply_chat_template(*args, **kwargs)
@property
def vocab(self):
"""
Returns the vocabulary as a list where the index corresponds to the token ID.
Returns:
List[str]: List of tokens in the vocabulary.
"""
id2vocab = {v: k for k, v in self.tokenizer.vocab.items()}
return [id2vocab[i] for i in range(len(id2vocab))]
@property
def pad_id(self):
"""
Gets the ID of the padding token.
Returns:
int or None: The ID of the padding token if it exists, None otherwise.
"""
if getattr(self, "pad_token") is None:
return None
return self.tokens_to_ids([getattr(self, "pad_token")])[0]
@property
def bos_id(self):
"""
Gets the ID of the beginning-of-sequence token.
Returns:
int or None: The ID of the BOS token if it exists, None otherwise.
"""
if getattr(self, "bos_token") is None:
return None
return self.tokens_to_ids([getattr(self, "bos_token")])[0]
@property
def eos_id(self):
"""
Gets the ID of the end-of-sequence token.
Returns:
int or None: The ID of the EOS token if it exists, None otherwise.
"""
if getattr(self, "eos_token") is None:
return None
return self.tokens_to_ids([getattr(self, "eos_token")])[0]
@property
def eod(self):
"""
Gets the ID of the end-of-document token (same as EOS token). Required for megatron-core compatibility.
Returns:
int: The ID of the EOD/EOS token.
"""
return self.tokens_to_ids([getattr(self, "eos_token")])[0]
@property
def sep_id(self):
"""
Gets the ID of the separator token.
Returns:
int or None: The ID of the separator token if it exists, None otherwise.
"""
if getattr(self, "sep_token") is None:
return None
return self.tokens_to_ids([getattr(self, "sep_token")])[0]
@property
def cls_id(self):
"""
Gets the ID of the classifier token.
Returns:
int or None: The ID of the classifier token if it exists, None otherwise.
"""
if getattr(self, "cls_token") is None:
return None
return self.tokens_to_ids([getattr(self, "cls_token")])[0]
@property
def unk_id(self):
"""
Gets the ID of the unknown token.
Returns:
int or None: The ID of the unknown token if it exists, None otherwise.
"""
if getattr(self, "unk_token") is None:
return None
return self.tokens_to_ids([getattr(self, "unk_token")])[0]
@property
def mask_id(self):
"""
Gets the ID of the mask token.
Returns:
int or None: The ID of the mask token if it exists, None otherwise.
"""
if getattr(self, "mask_token") is None:
return None
return self.tokens_to_ids([getattr(self, "mask_token")])[0]
@property
def name(self):
"""
Returns the name of the underlying HuggingFace tokenizer class.
Returns:
str: Name of the tokenizer class.
"""
return type(self.tokenizer).__name__
def save_pretrained(self, save_directory: str):
"""Saves tokenizer's vocabulary and other artifacts to the specified directory"""
return self.tokenizer.save_pretrained(save_directory)
vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py
+249
View File
@@ -0,0 +1,249 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Stage-2 Code2Wav for NemotronVoiceChat: RVQ-VAE codec decode to 22.05 kHz PCM.
``LLM_GENERATION`` stage consuming the talker's per-frame 31-quantizer code
stacks and decoding them with the vendored NeMo RVQ-VAE codec. The decode runs
in fp32 (the NeMo reference wraps codec decode in an fp32 context; bf16 collapses
the reconstruction).
Input contract (mirrors ``PersonaPlexCode2Wav``): the connector payload carries
``{"codes": {"audio": Tensor[frames, num_quantizers]}}`` per request; the
scheduler-side ``input_ids`` are placeholders. Codes must lie in
``[0, codebook_size)`` — out-of-range codes are a hard error (an early bf16-drift
signal upstream), never clamped.
"""
from __future__ import annotations
from collections.abc import Iterable, Mapping
from typing import Any
import torch
import torch.nn as nn
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm_omni.model_executor.models.output_templates import OmniOutput
logger = init_logger(__name__)
def validate_code_stack(codes: torch.Tensor, num_quantizers: int, codebook_size: int) -> torch.Tensor:
"""Validate and reshape a code payload to ``[frames, num_quantizers]``.
Out-of-range codes are a hard error (an early numeric-drift signal
upstream), never clamped.
"""
q = num_quantizers
if codes.dim() == 1:
if codes.numel() % q != 0:
raise ValueError(
f"NemotronVoiceChat code2wav received a flat code sequence of length "
f"{codes.numel()} which is not divisible by num_quantizers={q}."
)
codes = codes.reshape(-1, q)
if codes.dim() != 2 or codes.shape[-1] != q:
raise ValueError(f"NemotronVoiceChat code2wav expects codes shaped [frames, {q}], got {tuple(codes.shape)}.")
codes = codes.to(dtype=torch.long)
bad = (codes < 0) | (codes >= codebook_size)
if bool(bad.any()):
bad_vals = codes[bad]
raise ValueError(
"NemotronVoiceChat code2wav received out-of-range codec codes "
f"(min={int(bad_vals.min())}, max={int(bad_vals.max())}, count={int(bad.sum())}); "
f"valid range is [0, {codebook_size}). Refusing to clamp — this usually "
"signals numeric drift or a corrupted stage payload upstream."
)
return codes
class NemotronVoiceChatCode2Wav(nn.Module):
"""RVQ-VAE codec decoder stage (GenerationModelRunner)."""
input_modalities = "audio"
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
super().__init__()
self.vllm_config = vllm_config
self.config = vllm_config.model_config.hf_config
codec_cfg = getattr(self.config, "code2wav_config", None)
self._num_quantizers = int(getattr(codec_cfg, "num_quantizers", 31))
self._codebook_size = int(getattr(codec_cfg, "codebook_size", 1024))
self._sample_rate = int(getattr(codec_cfg, "sample_rate", 22050))
# Samples per 12.5 Hz frame (streaming chunks slice new audio by frame).
self._wav_per_frame = int(getattr(codec_cfg, "wav_to_token_ratio", 1764))
# Runner-facing capability flags (same contract as PersonaPlexCode2Wav).
self.have_multimodal_outputs = True
self.has_preprocess = False
self.has_postprocess = False
self.enable_update_additional_information = True
self.requires_raw_input_tokens = True
# Constructed in load_weights (owns the tts_model.audio_codec.* subtree).
self.audio_codec: nn.Module | None = None
# ------------------------------------------------------------------
# Runner-facing placeholder hooks.
# ------------------------------------------------------------------
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32)
def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None:
return None
def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, object]]:
return [{"meta": {"nemotron_voicechat_dummy_profile": True}} for _ in range(num_reqs)]
# ------------------------------------------------------------------
# Decode.
# ------------------------------------------------------------------
def _codes_from_runtime_info(self, runtime_info: Mapping[str, Any] | None) -> torch.Tensor | None:
"""Extract the codes payload; None means "no codes info at all".
An EMPTY tensor is returned as-is: it is an explicit empty payload
(e.g. zero acoustic frames after the prompt trim) and must yield empty
audio, NOT a fallback decode of the placeholder input_ids.
"""
if not isinstance(runtime_info, Mapping):
return None
if "codes" in runtime_info:
codes = runtime_info["codes"]
if isinstance(codes, Mapping):
codes = codes.get("audio")
else:
codes = runtime_info.get("codes.audio")
if isinstance(codes, list | tuple):
codes = torch.as_tensor(codes, dtype=torch.long)
if isinstance(codes, torch.Tensor):
return codes
return None
@torch.no_grad()
def forward(
self,
input_ids: torch.Tensor | None = None,
positions: torch.Tensor | None = None,
intermediate_tensors: Any = None,
inputs_embeds: torch.Tensor | None = None,
runtime_additional_information: list[dict[str, Any]] | None = None,
**kwargs: Any,
) -> OmniOutput:
sr_tensor = torch.tensor(self._sample_rate, dtype=torch.int32)
empty = torch.zeros((0,), dtype=torch.float32)
infos = runtime_additional_information or []
num_req = max(len(infos), 1)
audios: list[torch.Tensor] = [empty] * num_req
srs = [sr_tensor] * num_req
if self.audio_codec is None:
raise RuntimeError("NemotronVoiceChatCode2Wav.forward called before load_weights().")
device = next(self.audio_codec.parameters()).device
for i in range(num_req):
runtime_info = infos[i] if i < len(infos) else None
meta = runtime_info.get("meta") if isinstance(runtime_info, Mapping) else None
if isinstance(meta, Mapping) and meta.get("nemotron_voicechat_dummy_profile") is True:
continue
codes = self._codes_from_runtime_info(runtime_info)
if codes is None:
# No input_ids fallback (unlike PersonaPlexCode2Wav): this
# pipeline's scheduler-side prompt is all-zero placeholders, so
# decoding it would produce plausible-sounding garbage whenever
# n_frames % num_quantizers == 0. A missing payload is a bug.
raise ValueError(
"NemotronVoiceChat code2wav request carries no 'codes' payload in its "
"runtime info; the talker stage must ship the [frames, num_quantizers] "
"code stack (placeholder input_ids are never decoded)."
)
codes = validate_code_stack(codes, self._num_quantizers, self._codebook_size).to(device=device)
frames = codes.shape[0]
if frames == 0:
continue
# Async-chunk streaming: the payload carries the CUMULATIVE trimmed
# code history plus meta.left_context_size = frames whose audio was
# already emitted by earlier chunks. Mirroring NeMo's
# decode_one_audio_step (unbounded window), we re-decode the whole
# prefix and slice off only the new tail — seam-free by
# construction. Sync full-payload chunks carry no left context.
left_context_frames = 0
if isinstance(meta, Mapping) and meta.get("left_context_size") is not None:
left_context_frames = max(int(meta["left_context_size"]), 0)
elif isinstance(runtime_info, Mapping) and runtime_info.get("meta.left_context_size") is not None:
left_context_frames = max(int(runtime_info["meta.left_context_size"]), 0)
if left_context_frames >= frames:
continue # terminal empty chunk: all frames already emitted
lens = torch.tensor([frames], device=device, dtype=torch.long)
# NeMo decodes the codec strictly in fp32.
with torch.autocast(device_type=device.type, enabled=False):
wav, wav_len = self.audio_codec.decode(codes.unsqueeze(0), lens)
wav = wav.reshape(-1)[: int(wav_len.reshape(-1)[0])]
if left_context_frames > 0:
wav = wav[left_context_frames * self._wav_per_frame :]
audios[i] = wav.detach().to(dtype=torch.float32).cpu()
return OmniOutput(
text_hidden_states=None,
multimodal_outputs={"model_outputs": audios, "sr": srs},
)
def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput | tuple, **kwargs: Any) -> OmniOutput:
if isinstance(model_outputs, OmniOutput):
return model_outputs
raise TypeError(f"NemotronVoiceChatCode2Wav expected OmniOutput, got {type(model_outputs)}")
# ------------------------------------------------------------------
# Weights.
# ------------------------------------------------------------------
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
"""Build the vendored RVQ-VAE codec and load ``tts_model.audio_codec.*``.
The vLLM safetensors iterator streams the full unified checkpoint; only
the codec subtree is materialized here (everything else is skipped
without copying). The codec stays fp32.
"""
from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.ear_tts_vae_codec import (
RVQVAEModel,
)
codec_cfg = dict(getattr(self.config, "tts_cfg", {}).get("codec_config") or {})
if not codec_cfg:
raise ValueError(
"NemotronVoiceChat checkpoint config lacks "
"'model.speech_generation.model.codec_config'; cannot build the RVQ-VAE codec."
)
from omegaconf import DictConfig
codec = RVQVAEModel(DictConfig(codec_cfg))
prefix = "tts_model.audio_codec."
state: dict[str, torch.Tensor] = {}
for name, tensor in weights:
if name.startswith(prefix):
state[name.removeprefix(prefix)] = tensor.to(dtype=torch.float32)
if not state:
raise ValueError(
f"No '{prefix}*' tensors found in the checkpoint; the NemotronVoiceChat "
"unified model.safetensors is required for the code2wav stage."
)
missing, unexpected = codec.load_state_dict(state, strict=False)
if missing:
raise ValueError(f"NemotronVoiceChat code2wav is missing codec weights: {sorted(missing)[:10]} ...")
if unexpected:
logger.warning(
"NemotronVoiceChat code2wav ignoring %d unexpected codec tensors (e.g. %s)",
len(unexpected),
sorted(unexpected)[:5],
)
device = self.vllm_config.device_config.device
codec = codec.to(device=device, dtype=torch.float32).eval()
self.audio_codec = codec
logger.info(
"NemotronVoiceChat code2wav loaded RVQ-VAE codec: %d tensors, num_quantizers=%d, sr=%d",
len(state),
self._num_quantizers,
self._sample_rate,
)
return {name for name, _ in self.named_parameters()}
vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py
+490
View File
@@ -0,0 +1,490 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Stage-1 talker for NemotronVoiceChat: frame-locked text -> 31-quantizer codes.
``LLM_AR`` stage driven per frame by the vLLM engine, with the vendored NeMo
EAR-TTS model (28-layer Gemma3-style backbone + MoG head) executed inside the
per-request ``preprocess`` hook. Classifier-free guidance (internal batch
doubling) and MoG sampling happen inside the vendored forward exactly as in
NeMo — vLLM's sampler never touches the audio codes.
Execution contract (mirrors the NeMo ``NemotronVoiceChat.offline_inference``
TTS half):
* The request's ``additional_information`` carries the full frame-aligned text
timeline (built by ``thinker2talker_token_only``), INCLUDING the PAD prompt
region — the talker's KV state at the first acoustic frame depends on having
processed it (NeMo trims prompt-region codes only after the loop).
* vLLM prompt = one placeholder token (timeline position 0; NeMo's loop starts
at t=1). Each decode step t runs ``infer_codes_one_step`` with
``current_subword_id=timeline[t]``, ``prev_subword_id=timeline[t-1]`` (the
first step uses the warmup's last init subword, as in NeMo), and the previous
step's code stack; the resulting ``[1, 31]`` codes accumulate under
``("codes","audio")`` in the request state for the full-payload producer.
* Warmup: ``set_init_inputs(speaker_name=...)`` (pre-baked "Aria" latent) +
one ``tts_model(**init_inputs)`` forward seeds ``past_key_values`` and the
initial code, exactly as in the reference.
* Completion: after consuming timeline position ``T-1`` the stage emits a stop
placeholder token (the deploy yaml sets ``stop_token_ids: [1]``); codes for
positions ``1..T-1`` are shipped and the producer trims the prompt region.
Heavyweight per-request state (HF ``past_key_values`` etc.) lives in a
model-side session dict keyed by request id (PersonaPlex pattern), cleaned up
via ``on_requests_finished``.
"""
from __future__ import annotations
from collections.abc import Iterable
from typing import Any
import torch
import torch.nn as nn
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm_omni.model_executor.models.nemotron_voicechat.runtime_info import (
merge_runtime_info,
require_request_id,
)
from vllm_omni.model_executor.models.output_templates import OmniOutput
logger = init_logger(__name__)
_CONTINUE_TOKEN = 0
_STOP_TOKEN = 1
def sanitize_tts_model_cfg(tts_model_cfg: dict[str, Any]) -> dict[str, Any]:
"""Make the vendored DuplexEARTTS construction config-only.
Every parameter comes from the unified checkpoint streamed in
``load_weights``; the vendored constructors must never download or load
external pretrained weights. Returns a deep copy so the shared stage
config is never mutated:
* ``tts_config.pretrained_text_name`` would make ``RVQEARTTSModel``
download a full LLM backbone — dropped (the backbone builds from
``backbone_type``/``backbone_config``).
* ``pretrained_codec_model`` would make ``setup_audio_codec`` load an
external codec checkpoint at construction — dropped (the codec weights
arrive with the ``tts_model.audio_codec.*`` subtree, and the RVQ
embedding binding is overwritten by the same ``load_state_dict``).
* ``tts_config.cas_config.pretrained_tokenizer_name`` is dropped, matching
NeMo's own inference override (the CAS vocab is baked into the weights).
* ``pretrained_lm_name`` survives ONLY as the tokenizer reference (the
context-LM load path is disabled by ``context_hidden_size: null``); the
``NEMOTRON_VOICECHAT_LLM_PATH`` override keeps it air-gap friendly.
"""
import copy
import os
cfg = copy.deepcopy(tts_model_cfg)
cfg.pop("pretrained_codec_model", None)
tts_config = cfg.get("tts_config")
if isinstance(tts_config, dict):
tts_config.pop("pretrained_text_name", None)
cas_cfg = tts_config.get("cas_config")
if isinstance(cas_cfg, dict):
cas_cfg.pop("pretrained_tokenizer_name", None)
llm_override = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH")
if llm_override:
cfg["pretrained_lm_name"] = llm_override
return cfg
class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module):
"""vLLM AR stage wrapping the vendored NeMo EAR-TTS per-frame step."""
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
super().__init__()
self.vllm_config = vllm_config
self.config = vllm_config.model_config.hf_config
talker_cfg = getattr(self.config, "talker_config", None)
self._hidden = int(getattr(talker_cfg, "hidden_size", 1152))
self._vocab = max(int(getattr(talker_cfg, "vocab_size", 1024)), 2)
self._dtype = getattr(vllm_config.model_config, "dtype", torch.bfloat16)
self._speaker_name = getattr(self.config, "inference_speaker_name", "Aria") or "Aria"
# Omni AR runner contract flags.
self.have_multimodal_outputs = True
self.has_preprocess = True
self.has_postprocess = False
self.requires_full_prefix_cached_hidden_states = False
# Stateful per-request execution is only validated at batch size 1
# (the shipped offline scope); fail fast on silent multi-request use.
max_num_seqs = int(getattr(vllm_config.scheduler_config, "max_num_seqs", 1))
if max_num_seqs != 1:
raise NotImplementedError(
f"NemotronVoiceChat talker supports max_num_seqs=1 only (got {max_num_seqs}); "
"the vendored EAR-TTS session state is not batch-safe."
)
# Vendored DuplexEARTTS; constructed in load_weights.
self.tts: nn.Module | None = None
# request_id -> mutable session state (past_key_values, prev code, ...).
self._sessions: dict[str, dict[str, Any]] = {}
# ------------------------------------------------------------------
# vLLM model protocol (placeholder token loop).
# ------------------------------------------------------------------
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
return torch.zeros((input_ids.shape[0], self._hidden), device=input_ids.device, dtype=self._dtype)
def forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
intermediate_tensors: Any = None,
inputs_embeds: torch.Tensor | None = None,
**_: Any,
) -> torch.Tensor:
# The per-frame compute happens in preprocess; the "hidden state" only
# transports the stop flag (column 0) to compute_logits.
if inputs_embeds is not None:
return inputs_embeds
assert input_ids is not None
return self.embed_input_ids(input_ids)
def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> torch.Tensor:
if isinstance(hidden_states, OmniOutput):
hidden_states = hidden_states.text_hidden_states
hs = hidden_states
flag = hs[..., 0].to(torch.float32) # 1.0 => finished
logits = torch.full((hs.shape[0], self._vocab), -1.0e4, device=hs.device, dtype=torch.float32)
logits[:, _CONTINUE_TOKEN] = (1.0 - flag) * 10.0
logits[:, _STOP_TOKEN] = flag * 10.0
return logits
def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
if isinstance(model_outputs, OmniOutput):
return model_outputs
hidden = model_outputs
info_dicts = kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or []
codes_list: list[torch.Tensor] = []
prompt_len: int | None = None
for info in info_dicts:
if not isinstance(info, dict):
continue
codes = info.get("codes", {})
audio = codes.get("audio") if isinstance(codes, dict) else None
if isinstance(audio, torch.Tensor):
codes_list.append(audio)
meta = info.get("meta")
reported = meta.get("nvc_logical_prompt_len") if isinstance(meta, dict) else None
if reported is None:
reported = info.get("meta.nvc_logical_prompt_len")
if reported is not None:
prompt_len = int(reported)
if not codes_list:
return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
outputs: dict[str, Any] = {"codes": {"audio": torch.cat(codes_list, dim=0)}}
# The prompt length must ride the WIRE payload (multimodal_outputs):
# the async-chunk dispatcher hands exactly this dict to the
# talker2code2wav producer, while request.additional_information is
# overwritten concurrently by arriving thinker chunks.
if prompt_len is not None:
# 1-D on purpose: the wire-payload builder indexes element.shape[0].
outputs["meta"] = {"nvc_logical_prompt_len": torch.tensor([prompt_len], dtype=torch.long)}
return OmniOutput(
text_hidden_states=hidden,
multimodal_outputs=outputs,
)
# ------------------------------------------------------------------
# Per-frame vendored TTS step.
# ------------------------------------------------------------------
def _try_timeline(self, info: dict[str, Any]) -> torch.Tensor | None:
"""Extract the frame-locked text timeline from either payload shape.
Sync mode ships ``nvc_text_timeline`` (+ ``nvc_logical_prompt_len``) in
additional_information; async-chunk mode replaces the request info with
the connector payload each chunk, carrying the CUMULATIVE timeline under
``ids.all`` (and the logical prompt length under
``meta.num_processed_tokens``). Returns None when neither is present.
"""
timeline = info.get("nvc_text_timeline")
if timeline is not None and info.get("nvc_logical_prompt_len") is None:
raise ValueError(
"NemotronVoiceChat talker got 'nvc_text_timeline' without "
"'nvc_logical_prompt_len'; the code2wav producer needs the prompt "
"length to trim prompt-region codes."
)
if timeline is None:
ids = info.get("ids")
timeline = ids.get("all") if isinstance(ids, dict) else None
if timeline is None:
timeline = info.get("ids.all")
if timeline is None:
return None
if isinstance(timeline, torch.Tensor):
return timeline.reshape(-1).to(torch.long)
return torch.as_tensor(list(timeline), dtype=torch.long)
@staticmethod
def _logical_prompt_len(info: dict[str, Any]) -> int | None:
"""Logical prompt length from either payload shape (sync/async)."""
value = info.get("nvc_logical_prompt_len")
if value is None:
meta = info.get("meta")
value = meta.get("num_processed_tokens") if isinstance(meta, dict) else None
if value is None:
value = info.get("meta.num_processed_tokens")
return int(value) if value is not None else None
@staticmethod
def _upstream_finished(info: dict[str, Any]) -> bool:
"""True when the async producer marked its final chunk (meta.finished)."""
meta = info.get("meta")
flag = meta.get("finished") if isinstance(meta, dict) else None
if flag is None:
flag = info.get("meta.finished")
if isinstance(flag, torch.Tensor):
return bool(flag.reshape(-1)[0].item()) if flag.numel() else False
return bool(flag)
def _timeline(self, info: dict[str, Any]) -> torch.Tensor:
timeline = self._try_timeline(info)
if timeline is None:
# Explicit failure — no prompt_token_ids fallback: an implicit
# timeline would silently skip the prompt-region trim downstream.
raise ValueError(
"NemotronVoiceChat talker request is missing its timeline metadata "
"('nvc_text_timeline' in additional_information, or 'ids.all' on the "
"async-chunk payload); the stage cannot synthesize speech without "
"the thinker's frame-locked tokens."
)
return timeline
def _init_session(self, request_id: str, info: dict[str, Any], device: torch.device) -> dict[str, Any]:
assert self.tts is not None
timeline = self._timeline(info).to(device)
if timeline.numel() < 2:
raise ValueError(
f"NemotronVoiceChat talker timeline for request {request_id} has "
f"{timeline.numel()} positions; need at least 2 (prompt + one frame)."
)
# Warmup: pre-baked speaker latent + init prompt forward (NeMo
# offline_inference lines around get_init_inputs/first forward).
self.tts.set_init_inputs(speaker_name=self._speaker_name)
init_inputs = self.tts.get_init_inputs(B=1)
generation_config = self.tts._get_generation_config(self._guidance_enabled())
init_inputs.update({"use_cache": True, "past_key_values": None, "guidance_enabled": self._guidance_enabled()})
with torch.inference_mode():
outputs = self.tts.tts_model(**init_inputs)
prompt_len = self._logical_prompt_len(info)
if prompt_len is None:
raise ValueError(
"NemotronVoiceChat talker request is missing its logical prompt length "
"('nvc_logical_prompt_len' in additional_information, or async "
"meta.num_processed_tokens); the code2wav producer cannot trim the "
"prompt-region codes without it."
)
session = {
"timeline": timeline,
"prompt_len": prompt_len,
"step": 1, # NeMo's loop starts at t=1
"past_key_values": outputs.past_key_values,
"code": init_inputs["code"][:, -1:],
"first_context_subword_id": init_inputs["subword_ids"][:, -1].unsqueeze(-1),
"tts_initialized": False,
"generation_config": generation_config,
"codes_rows": [],
# Sync mode ships the whole timeline up front via
# nvc_text_timeline; async-chunk mode grows it across chunks and
# the last chunk sets meta.finished.
"sync_mode": info.get("nvc_text_timeline") is not None,
"upstream_finished": self._upstream_finished(info) or info.get("nvc_text_timeline") is not None,
}
self._sessions[request_id] = session
return session
def _refresh_session_timeline(self, session: dict[str, Any], info: dict[str, Any], device: torch.device) -> None:
"""Adopt a longer timeline from the latest async-chunk payload."""
latest = self._try_timeline(info)
if latest is not None and latest.numel() > int(session["timeline"].numel()):
session["timeline"] = latest.to(device)
if self._upstream_finished(info):
session["upstream_finished"] = True
def _guidance_enabled(self) -> bool:
return bool(getattr(self.config, "tts_cfg", {}).get("inference_guidance_enabled", True))
def _step_session(self, session: dict[str, Any]) -> tuple[torch.Tensor, bool]:
"""One NeMo TTS step; returns (codes_row [1, Q], finished)."""
assert self.tts is not None
timeline: torch.Tensor = session["timeline"]
t = int(session["step"])
total = int(timeline.numel())
if t >= total:
# The drain loop in preprocess never enters here; a direct call
# past the received timeline is a programming error. Hard-fail
# rather than fabricating a frame — a repeated/guessed subword
# would silently desync the audio from the text channel.
raise RuntimeError(
f"NemotronVoiceChat talker stepped past the received timeline: step {t} "
f"needs timeline position {t} but only {total} positions have arrived "
f"(upstream_finished={session['upstream_finished']})."
)
current_subword_id = timeline[t].reshape(1, 1)
if not session["tts_initialized"]:
prev_subword_id = session["first_context_subword_id"]
session["tts_initialized"] = True
else:
prev_subword_id = timeline[t - 1].reshape(1, 1)
current_subword_mask = torch.ones(1, 1, device=timeline.device, dtype=torch.bool)
with torch.inference_mode():
code, past_key_values = self.tts.infer_codes_one_step(
current_subword_id=current_subword_id,
prev_subword_id=prev_subword_id,
current_subword_mask=current_subword_mask,
prev_audio_tokens=session["code"],
past_key_values=session["past_key_values"],
guidance_enabled=self._guidance_enabled(),
generation_config=session["generation_config"],
ignore_eos_flag_stop=True,
)
session["code"] = code
session["past_key_values"] = past_key_values
session["step"] = t + 1
# Async mode: exhausting the received timeline only ends the request
# once the upstream marked its final chunk; otherwise more entries are
# in flight and the scheduler parks us until they arrive.
finished = (t + 1) >= int(session["timeline"].numel()) and bool(session["upstream_finished"])
return code.reshape(1, -1).to(torch.long), finished
def preprocess(
self,
input_ids: torch.Tensor,
input_embeds: torch.Tensor | None,
**info_dict: Any,
) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
info = merge_runtime_info(info_dict)
device = input_ids.device
span = int(input_ids.shape[0])
is_prefill_raw = info.get("_omni_is_prefill")
is_prefill = bool(is_prefill_raw) if isinstance(is_prefill_raw, bool) else span > 1
request_id = require_request_id(info, "talker")
embeds = torch.zeros((span, self._hidden), device=device, dtype=self._dtype)
if is_prefill or request_id not in self._sessions:
# The 1-token vLLM prefill is the session boundary: run the vendored
# warmup here so the first decode step is a pure NeMo t=1 step.
self._init_session(request_id, info, device)
return input_ids, embeds, {}
session = self._sessions[request_id]
self._refresh_session_timeline(session, info, device)
# Sync mode: one NeMo TTS step per engine step (the verified parity
# path, unchanged). Async-chunk mode: DRAIN every received timeline
# position on each wake. The engine tokens are placeholders
# (CONTINUE/STOP) and the codes ride the cumulative payload, so
# multiple TTS steps per engine step are safe — and necessary:
# * the save_async background thread can coalesce thinker steps into
# one chunk (fewer wakes than new positions), and
# * the prompt-region PAD steps (logical_prompt_len - 1 of them) all
# become available with chunk 0 and must not cost one upstream
# chunk each.
finished = False
steps_run = 0
while int(session["step"]) < int(session["timeline"].numel()):
codes_row, finished = self._step_session(session)
session["codes_rows"].append(codes_row.cpu())
steps_run += 1
if session["sync_mode"]:
break
if steps_run == 0:
# Zero-progress wake (duplicate/terminal-marker chunk): no TTS step
# was run, so emit a plain CONTINUE placeholder — never a
# fabricated frame. The request finishes only when the upstream is
# exhausted AND the timeline is fully consumed.
finished = bool(session["upstream_finished"])
if finished:
embeds[:, 0] = 1.0
return input_ids, embeds, {}
cumulative = torch.cat(session["codes_rows"], dim=0)
if finished:
embeds[:, 0] = 1.0 # stop flag -> compute_logits emits the stop token
info_update: dict[str, Any] = {
"codes": {"audio": cumulative},
# nvc_logical_prompt_len rides EVERY step's payload: the request's
# additional_information is overwritten both by arriving thinker
# chunks and by these model updates, so the async code2wav producer
# reads the prompt length from the per-step multimodal_output
# instead of racing the request-level dict.
"meta": {"nvc_tts_step": int(session["step"]), "nvc_logical_prompt_len": int(session["prompt_len"])},
}
return input_ids, embeds, info_update
def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
for request_id in finished_req_ids:
self._sessions.pop(str(request_id), None)
# ------------------------------------------------------------------
# Weights.
# ------------------------------------------------------------------
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
"""Build the vendored DuplexEARTTS and load the ``tts_model.*`` subtree.
The backbone/MoG run in the engine dtype (bf16 by default); the codec
submodule inside DuplexEARTTS stays fp32 (it is not used by this stage —
decode happens in code2wav — but keeping its dtype canonical avoids
surprises if warmup paths touch it).
"""
from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.duplex_ear_tts import (
DuplexEARTTS,
)
tts_model_cfg = dict(getattr(self.config, "tts_cfg", {}) or {})
if not tts_model_cfg:
raise ValueError(
"NemotronVoiceChat checkpoint config lacks 'model.speech_generation.model'; cannot build the talker."
)
tts_model_cfg = sanitize_tts_model_cfg(tts_model_cfg)
tts_data = dict(getattr(self.config, "tts_data", {}) or {})
tts_data.setdefault("source_sample_rate", 22050)
tts_data.setdefault("target_sample_rate", int(getattr(self.config, "target_sample_rate", 22050)))
tts_data.setdefault("frame_length", float(getattr(self.config, "frame_length", 0.08)))
# DuplexEARTTS consumes the whole speech_generation section layout
# ({"data": ..., "model": ...}), matching NeMo's constructor call.
tts = DuplexEARTTS({"data": tts_data, "model": tts_model_cfg})
prefix = "tts_model."
state: dict[str, torch.Tensor] = {}
for name, tensor in weights:
if not name.startswith(prefix):
continue
target = name.removeprefix(prefix)
if target.startswith("audio_codec."):
state[target] = tensor.to(dtype=torch.float32)
else:
state[target] = tensor.to(dtype=self._dtype)
if not state:
raise ValueError(
"No 'tts_model.*' tensors found in the checkpoint; the NemotronVoiceChat "
"unified model.safetensors is required for the talker stage."
)
missing, unexpected = tts.load_state_dict(state, strict=False)
real_missing = [m for m in missing if not m.startswith("audio_codec.")]
if real_missing:
raise ValueError(f"NemotronVoiceChat talker is missing TTS weights: {sorted(real_missing)[:10]} ...")
if unexpected:
logger.warning(
"NemotronVoiceChat talker ignoring %d unexpected TTS tensors (e.g. %s)",
len(unexpected),
sorted(unexpected)[:5],
)
device = self.vllm_config.device_config.device
tts = tts.to(device=device)
tts.tts_model.to(dtype=self._dtype)
if hasattr(tts, "audio_codec") and tts.audio_codec is not None:
tts.audio_codec.to(dtype=torch.float32)
self.tts = tts.eval()
logger.info(
"NemotronVoiceChat talker loaded EAR-TTS: %d tensors, speaker=%s, guidance=%s",
len(state),
self._speaker_name,
self._guidance_enabled(),
)
return {name for name, _ in self.named_parameters()}
vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py
+596
View File
@@ -0,0 +1,596 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Stage-0 thinker for NemotronVoiceChat: 16 kHz speech -> frame-locked text.
``LLM_AR`` stage combining:
* the vendored NeMo FastConformer perception (mel + 24-layer encoder + linear
projection to the LLM width, 12.5 Hz frames);
* the additive channel fusion from the NeMo duplex STT model
(``fused[t] = E(text_prev) * w_text + timeline[t] * w_audio +
E(func_prev) * w_func`` with the checkpoint's AddFusion weights);
* the NemotronH (hybrid Mamba) backbone on vLLM paged execution via
``init_vllm_registered_model`` — this is where vLLM earns its keep;
* the separate ``lm_head`` (text channel, sampled by the vLLM engine) and
``function_head`` (greedy, maintained internally per request).
Timeline contract (verified against NeMo ``_init_inference``/``_step_zero``/
``_step_inference``; the golden shape test guards it):
* The vLLM prompt covers ``logical_prompt_token_len + 1`` positions: the system
prompt tokens (``[bos] + tokens + [eos]``) fused with a PAD text channel
(NeMo's ``_get_bos_embedding`` returns the PAD embedding, so position 0 is
uniform with the rest of the prompt region), plus acoustic frame 0 as the
final prefill position. The first engine-sampled token is therefore NeMo's
``gen_text[prompt_len]`` — the first acoustic-frame prediction.
* Each decode step t consumes conformer frame t (held in per-request state)
fused with the previously sampled text token and the previous greedy
function token. Text EOS never terminates generation (``ignore_eos``); the
example sets ``max_tokens = acoustic_frame_count``. Never set ``min_tokens``:
the tokenizer's eos_token is <SPECIAL_12> — the frame-locked PAD/silence
token — and min_tokens masks "EOS" to -inf, forbidding silence entirely.
* The function channel exists for numerical correctness (fusion weight 2.0)
even though tool calling is out of scope; the greedy token timeline is
exposed for debugging via the request info.
"""
from __future__ import annotations
import os
from collections.abc import Iterable
from typing import Any
import numpy as np
import torch
import torch.nn as nn
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.model_executor.models.interfaces import HasInnerState, IsHybrid
from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
from vllm_omni.model_executor.models.nemotron_voicechat.runtime_info import (
merge_runtime_info,
require_request_id,
)
from vllm_omni.model_executor.models.output_templates import OmniOutput
logger = init_logger(__name__)
_LM_ARCHITECTURE = "NemotronHForCausalLM"
def validate_timeline_fits_max_model_len(
logical_prompt_token_len: int,
acoustic_frame_count: int,
max_model_len: int,
) -> None:
"""Reject audio whose frame-locked timeline exceeds the stage limit.
The vLLM request occupies ``logical_prompt_token_len + 1`` prefill
positions (the +1 is the acoustic-frame-0 placeholder) plus
``acoustic_frame_count`` sampled tokens, so the binding sequence size is
one MORE than NeMo's logical timeline (``prompt + frames``).
"""
logical_total_len = logical_prompt_token_len + acoustic_frame_count
vllm_prefill_len = logical_prompt_token_len + 1
vllm_total_len = vllm_prefill_len + acoustic_frame_count
if vllm_total_len > max_model_len:
raise ValueError(
"NemotronVoiceChat input exceeds the thinker stage's max_model_len: "
f"vllm_prefill_len={vllm_prefill_len} + "
f"acoustic_frame_count={acoustic_frame_count} = "
f"vllm_total_len={vllm_total_len} > max_model_len={max_model_len} "
f"(logical_prompt_token_len={logical_prompt_token_len}, "
f"logical_total_len={logical_total_len}). Use shorter audio or raise the "
"stage's max_model_len in the deploy yaml."
)
def compute_acoustic_frame_count(nemo_stt_cfg: dict[str, Any], num_samples: int) -> int:
"""Exact 12.5 Hz frame count for ``num_samples`` of 16 kHz audio.
Replicates the vendored mel featurizer length arithmetic followed by the
dw-striding conv subsampling length recurrence (NOT ``duration / 0.08``).
The thinker validates this against the real perception output at prefill,
and the golden shape test pins it against the NeMo reference.
"""
from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.asr.subsampling import (
calc_length,
)
perception = nemo_stt_cfg.get("perception") or {}
pre = perception.get("preprocessor") or {}
encoder = perception.get("encoder") or {}
sample_rate = int(pre.get("sample_rate", 16000))
window_stride = float(pre.get("window_stride", 0.01))
window_size = float(pre.get("window_size", 0.025))
n_fft = int(pre.get("n_fft") or 2 ** int(np.ceil(np.log2(window_size * sample_rate))))
hop = int(window_stride * sample_rate)
pad_amount = n_fft // 2 * 2 # torch.stft center padding on both sides
mel_len = (num_samples + pad_amount - n_fft) // hop + 1
subsampling_factor = int(encoder.get("subsampling", "dw_striding") and encoder.get("subsampling_factor", 8))
conv_kernel = int(encoder.get("subsampling_conv_kernel_size", 3) or 3)
stride = 2
repeat = int(np.log2(subsampling_factor))
# Mirror the vendored ConvSubsampling padding exactly: causal downsampling
# (this checkpoint sets encoder.causal_downsampling=true) pads
# left=kernel-1 plus right=stride-1, i.e. one more column than the
# non-causal symmetric 2*(kernel//2) — which shifts the output length by
# one for roughly half of all input durations.
if bool(encoder.get("causal_downsampling", False)):
all_paddings = (conv_kernel - 1) + (stride - 1)
else:
all_paddings = 2 * (conv_kernel // 2)
out = calc_length(
torch.tensor([mel_len]),
all_paddings=all_paddings,
kernel_size=conv_kernel,
stride=stride,
ceil_mode=False,
repeat_num=repeat,
)
return int(out[0])
class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState, IsHybrid):
"""vLLM-native thinker (perception + fusion + NemotronH + dual heads).
Declares the hybrid-Mamba interfaces (``IsHybrid``/``HasInnerState``) so
vLLM configures the mamba state cache for the wrapped NemotronH backbone;
the state shape/dtype classmethods mirror ``NemotronHForCausalLM`` but read
the NemotronH sub-config (``get_text_config()``) instead of the top-level
checkpoint config.
"""
@classmethod
def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig) -> tuple[torch.dtype, torch.dtype]:
from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM
return NemotronHForCausalLM.get_mamba_state_dtype_from_config(vllm_config)
@classmethod
def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig) -> tuple[tuple[int, int], tuple[int, int, int]]:
from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator
parallel_config = vllm_config.parallel_config
hf_config = vllm_config.model_config.hf_config.get_text_config()
intermediate_size = hf_config.mamba_num_heads * hf_config.mamba_head_dim
return MambaStateShapeCalculator.mamba2_state_shape(
intermediate_size=intermediate_size,
tp_world_size=parallel_config.tensor_parallel_size,
n_groups=hf_config.n_groups,
num_heads=hf_config.mamba_num_heads,
head_dim=hf_config.mamba_head_dim,
state_size=hf_config.ssm_state_size,
conv_kernel=hf_config.conv_kernel,
num_spec=vllm_config.num_speculative_tokens,
)
@classmethod
def get_mamba_state_copy_func(cls):
from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM
return NemotronHForCausalLM.get_mamba_state_copy_func()
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
super().__init__()
self.vllm_config = vllm_config
config = vllm_config.model_config.hf_config
self.config = config
stt_cfg: dict[str, Any] = dict(getattr(config, "stt_cfg", {}) or {})
if not stt_cfg:
raise ValueError("NemotronVoiceChat checkpoint config lacks 'model.stt.model'; cannot build the thinker.")
self.stt_cfg = stt_cfg
text_cfg = config.get_text_config()
self._hidden = int(getattr(text_cfg, "hidden_size", 4480))
self._vocab = int(getattr(text_cfg, "vocab_size", 131072))
self._dtype = getattr(vllm_config.model_config, "dtype", torch.bfloat16)
# NemotronH backbone (+ lm_head) on vLLM native layers.
self.language_model = init_vllm_registered_model(
vllm_config=vllm_config,
hf_config=text_cfg,
prefix=maybe_prefix(prefix, "language_model"),
architectures=[_LM_ARCHITECTURE],
)
# Shared channel embedding + separate function head (both distinct
# checkpoints subtrees from the backbone).
self.embed_tokens = nn.Embedding(self._vocab, self._hidden)
self.function_head = nn.Linear(self._hidden, self._vocab, bias=False)
# Vendored perception (mel + FastConformer + projection).
from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.perception import (
AudioPerceptionModule,
)
perception_cfg = stt_cfg.get("perception")
if not perception_cfg:
raise ValueError("'model.stt.model.perception' missing from the checkpoint config.")
from omegaconf import DictConfig
self.perception = AudioPerceptionModule(DictConfig(perception_cfg))
# AddFusion weights (the checkpoint uses fuse_method="add" defaults;
# the same duplex_*_weight keys double as the fusion weights).
use_function_head = bool(stt_cfg.get("use_function_head", True))
self._w_text = float(stt_cfg.get("duplex_text_channel_weight", 1.0))
self._w_audio = float(stt_cfg.get("duplex_user_channel_weight", 1.0))
self._w_func = float(stt_cfg.get("duplex_function_channel_weight", 1.0)) if use_function_head else 0.0
self._use_function_head = use_function_head
fuse_method = stt_cfg.get("fuse_method", "add") or "add"
if fuse_method != "add":
raise NotImplementedError(
f"NemotronVoiceChat thinker only implements fuse_method='add' "
f"(this checkpoint uses it); got {fuse_method!r}."
)
if bool(stt_cfg.get("predict_user_text", False)):
raise NotImplementedError(
"NemotronVoiceChat thinker does not implement the user-ASR channel "
"(predict_user_text=True); this checkpoint has it disabled."
)
self._special_ids: dict[str, int] | None = None
# Omni AR runner contract.
self.have_multimodal_outputs = True
self.has_preprocess = True
self.has_postprocess = True # per-step greedy function-channel token
self.requires_full_prefix_cached_hidden_states = False
# The per-step function token is a scalar; no GPU-resident buffers needed.
self.gpu_resident_buffer_keys: set[tuple[str, str]] = set()
# Stateful per-request execution is only validated at batch size 1
# (the shipped offline scope); fail fast on silent multi-request use.
max_num_seqs = int(getattr(vllm_config.scheduler_config, "max_num_seqs", 1))
if max_num_seqs != 1:
raise NotImplementedError(
f"NemotronVoiceChat thinker supports max_num_seqs=1 only (got {max_num_seqs}); "
"multi-request batching of the per-request fusion state is not implemented."
)
# request_id -> per-request state (audio frames, prefill embeds, counters).
self._sessions: dict[str, dict[str, Any]] = {}
# ------------------------------------------------------------------
# Special token ids (from the checkpoint's token strings).
# ------------------------------------------------------------------
def _resolve_special_ids(self) -> dict[str, int]:
if self._special_ids is not None:
return self._special_ids
from transformers import AutoTokenizer
ref = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or self.stt_cfg.get(
"pretrained_llm", "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
)
tok = AutoTokenizer.from_pretrained(ref, trust_remote_code=False)
ids = {
"bos": tok.convert_tokens_to_ids(self.stt_cfg.get("bos_token", "<s>")),
"eos": tok.convert_tokens_to_ids(self.stt_cfg.get("eos_token", "</s>")),
"pad": tok.convert_tokens_to_ids(self.stt_cfg.get("pad_token", "<SPECIAL_12>")),
}
for name, value in ids.items():
if value is None or value < 0:
raise ValueError(f"NemotronVoiceChat could not resolve the {name} token id from {ref!r}.")
self._special_ids = ids
return ids
# ------------------------------------------------------------------
# vLLM model protocol.
# ------------------------------------------------------------------
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
# Placeholder; preprocess replaces every position with fused embeds.
return torch.zeros((input_ids.shape[0], self._hidden), device=input_ids.device, dtype=self._dtype)
def forward(
self,
input_ids: torch.Tensor | None,
positions: torch.Tensor,
intermediate_tensors: Any = None,
inputs_embeds: torch.Tensor | None = None,
**_: Any,
) -> torch.Tensor:
return self.language_model(
input_ids=None,
positions=positions,
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds,
)
def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> torch.Tensor:
if isinstance(hidden_states, OmniOutput):
hidden_states = hidden_states.text_hidden_states
# vLLM >= 0.25 dropped the sampling_metadata parameter.
return self.language_model.compute_logits(hidden_states)
def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
if isinstance(model_outputs, OmniOutput):
return model_outputs
return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs={})
# ------------------------------------------------------------------
# Fusion helpers.
# ------------------------------------------------------------------
def _fuse(
self,
text_ids: torch.Tensor,
timeline_embeds: torch.Tensor,
func_ids: torch.Tensor,
) -> torch.Tensor:
"""AddFusion over the three channels; shapes [T] ids, [T, H] timeline."""
text_e = self.embed_tokens(text_ids)
fused = text_e * self._w_text + timeline_embeds.to(text_e.dtype) * self._w_audio
if self._use_function_head:
fused = fused + self.embed_tokens(func_ids) * self._w_func
return fused
# ------------------------------------------------------------------
# Per-request session.
# ------------------------------------------------------------------
def _load_audio(self, info: dict[str, Any], device: torch.device) -> torch.Tensor:
audio = info.get("nvc_audio")
if audio is None:
raise ValueError(
"NemotronVoiceChat thinker request is missing its audio: pass "
"additional_information={'nvc_audio': <16 kHz mono float waveform>, "
"'nvc_sr': 16000, 'nvc_prompt_token_ids': [...]}."
)
wav = torch.as_tensor(np.asarray(audio), dtype=torch.float32, device=device).reshape(-1)
if wav.numel() == 0:
raise ValueError("NemotronVoiceChat thinker received empty audio input.")
sr = int(info.get("nvc_sr", 16000))
expected_sr = int(getattr(self.config, "source_sample_rate", 16000))
if sr != expected_sr:
raise ValueError(
f"NemotronVoiceChat thinker expects {expected_sr} Hz audio, got {sr} Hz — "
"resample client-side (the offline example does this)."
)
return wav
def _init_session(self, request_id: str, info: dict[str, Any], device: torch.device) -> dict[str, Any]:
pad_id = self._resolve_special_ids()["pad"]
wav = self._load_audio(info, device)
prompt_ids = info.get("nvc_prompt_token_ids")
if prompt_ids is None:
raise ValueError(
"NemotronVoiceChat thinker request is missing 'nvc_prompt_token_ids' "
"(the system prompt token ids, [bos] + tokens + [eos])."
)
prompt_ids = torch.as_tensor(list(prompt_ids), dtype=torch.long, device=device)
if prompt_ids.numel() == 0:
raise ValueError("NemotronVoiceChat thinker requires a non-empty system prompt.")
# Perception: mel + FastConformer + projection -> [N, hidden] frames.
with torch.inference_mode():
encoded, encoded_len = self.perception(
input_signal=wav.unsqueeze(0).to(next(self.perception.parameters()).dtype),
input_signal_length=torch.tensor([wav.numel()], device=device),
)
n_frames = int(encoded_len.reshape(-1)[0])
audio_embeds = encoded[0, :n_frames].to(self._dtype) # [N, H]
expected = info.get("nvc_expected_frames")
if expected is not None and int(expected) != n_frames:
raise ValueError(
f"NemotronVoiceChat frame-count mismatch: the request was sized for "
f"{int(expected)} acoustic frames but perception produced {n_frames}. "
"The example must size max_tokens with compute_acoustic_frame_count() (never min_tokens)."
)
p = int(prompt_ids.numel())
validate_timeline_fits_max_model_len(
logical_prompt_token_len=p,
acoustic_frame_count=n_frames,
max_model_len=int(self.vllm_config.model_config.max_model_len),
)
# Fused prefill embeds for positions 0..P (P prompt slots + frame 0).
# All prompt positions use the PAD text channel (NeMo's
# _get_bos_embedding returns the PAD embedding) and PAD function channel.
timeline = torch.cat([self.embed_tokens(prompt_ids).to(self._dtype), audio_embeds[0:1]], dim=0) # [P+1, H]
pad_ids = torch.full((p + 1,), pad_id, dtype=torch.long, device=device)
prefill = self._fuse(pad_ids, timeline, pad_ids) # [P+1, H]
session = {
"audio_embeds": audio_embeds,
"prefill_embeds": prefill,
"prompt_len": p,
"n_frames": n_frames,
"func_token": int(pad_id),
"decode_step": 0,
}
self._sessions[request_id] = session
return session
# ------------------------------------------------------------------
# Omni AR hooks.
# ------------------------------------------------------------------
def preprocess(
self,
input_ids: torch.Tensor,
input_embeds: torch.Tensor | None,
**info_dict: Any,
) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
info = merge_runtime_info(info_dict)
device = input_ids.device
span = int(input_ids.shape[0])
is_prefill_raw = info.get("_omni_is_prefill")
is_prefill = bool(is_prefill_raw) if isinstance(is_prefill_raw, bool) else span > 1
request_id = require_request_id(info, "thinker")
if is_prefill or request_id not in self._sessions:
session = self._sessions.get(request_id) or self._init_session(request_id, info, device)
offset = max(0, int(info.get("_omni_num_computed_tokens", 0) or 0))
prefill: torch.Tensor = session["prefill_embeds"]
if offset + span > prefill.shape[0]:
raise ValueError(
f"NemotronVoiceChat thinker prefill span [{offset}, {offset + span}) exceeds "
f"the fused prompt length {prefill.shape[0]}; the request prompt must be "
"the system prompt tokens plus exactly one placeholder position."
)
return input_ids, prefill[offset : offset + span].to(self._dtype), {}
session = self._sessions[request_id]
pad_id = self._resolve_special_ids()["pad"]
# Function channel: greedy argmax over the PREVIOUS position's hidden,
# computed by postprocess right after each forward. At the first decode
# step this is the prefill's last position (acoustic frame 0), matching
# NeMo's gen_function[prompt_len]. The cumulative debug timeline
# (nvc_function_tokens) is owned by postprocess, which also covers the
# final acoustic frame (no later preprocess exists for it).
if self._use_function_head:
prev_func = info.get("nvc_prev_function_token")
if prev_func is not None:
token = int(prev_func.reshape(-1)[0]) if isinstance(prev_func, torch.Tensor) else int(prev_func)
session["func_token"] = token
# Position-addressed frame index: decode of vLLM position P+k consumes
# acoustic frame k (frame 0 was the last prefill position). Deriving it
# from _omni_num_computed_tokens keeps preprocess idempotent — if the
# runner ever re-enters the same decode position (preemption or
# recompute), the audio pointer stays aligned with the text channel
# instead of silently advancing. Falls back to a call counter only if
# the runner did not inject the key.
computed = info.get("_omni_num_computed_tokens")
if computed is not None:
frame_idx = int(computed) - int(session["prompt_len"])
else:
frame_idx = int(session["decode_step"]) + 1
session["decode_step"] = frame_idx
n_frames = int(session["n_frames"])
if frame_idx < 1:
raise ValueError(
f"NemotronVoiceChat thinker derived decode frame index {frame_idx} "
f"(computed_tokens={computed}, prompt_len={session['prompt_len']}); "
"decode must start at acoustic frame 1."
)
if frame_idx >= n_frames:
raise ValueError(
f"NemotronVoiceChat thinker was asked for decode step {frame_idx} but only "
f"{n_frames} acoustic frames exist; max_tokens must equal the frame count."
)
text_prev = input_ids.reshape(-1)[:1].to(torch.long)
func_prev = torch.full_like(text_prev, session["func_token"] if self._use_function_head else pad_id)
frame = session["audio_embeds"][frame_idx : frame_idx + 1]
fused = self._fuse(text_prev, frame, func_prev) # [1, H]
info_update: dict[str, Any] = {
"meta": {"nvc_frame": frame_idx},
"nvc_text_pad_id": pad_id,
}
return input_ids, fused, info_update
def postprocess(self, hidden_states: torch.Tensor, **kwargs: Any) -> dict[str, Any]:
"""Greedy function-channel token from this step's last hidden.
Runs after EVERY forward (prefill and each decode step), so the final
acoustic frame's token is covered too — NeMo produces gen_function for
every non-prompt position. ``nvc_prev_function_token`` is load-bearing
(it feeds the next step's fusion at weight 2.0). The cumulative debug
timeline (``nvc_function_tokens``, one entry per acoustic frame) costs
a host sync plus a growing concat per step, so it is only accumulated
when ``NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE=1``.
On multi-chunk prefill, only the final chunk (whose forward covers the
last prefill position, acoustic frame 0) contributes a token.
"""
if hidden_states is None or hidden_states.numel() == 0 or not self._use_function_head:
return {}
if kwargs.get("_omni_is_prefill"):
computed = int(kwargs.get("_omni_num_computed_tokens", 0) or 0)
prompt_len = int(kwargs.get("_omni_prompt_len", 0) or 0)
if prompt_len and computed + int(hidden_states.shape[0]) < prompt_len:
return {} # intermediate prefill chunk: frame 0 not reached yet
with torch.inference_mode():
func_logits = self.function_head(hidden_states[-1, :].reshape(1, -1).to(self._dtype))
token = func_logits.argmax(dim=-1).reshape(()).detach()
update: dict[str, Any] = {"nvc_prev_function_token": token}
if os.environ.get("NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", "0") == "1":
existing = kwargs.get("nvc_function_tokens")
if isinstance(existing, torch.Tensor) and existing.numel() > 0:
update["nvc_function_tokens"] = torch.cat([existing.reshape(-1).cpu(), token.reshape(1).cpu()])
else:
update["nvc_function_tokens"] = token.reshape(1).cpu()
return update
def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
for request_id in finished_req_ids:
self._sessions.pop(str(request_id), None)
# ------------------------------------------------------------------
# Weights.
# ------------------------------------------------------------------
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
"""Route the unified checkpoint into the thinker's components.
* ``stt_model.llm.*`` -> vLLM NemotronH (renamed to ``backbone.*``; the
separate ``stt_model.embed_tokens`` doubles as ``backbone.embeddings``
since the checkpoint keeps a single shared text embedding).
* ``stt_model.lm_head.weight`` -> the backbone's ``lm_head``.
* ``stt_model.embed_tokens/function_head`` -> the fusion embedding /
function head.
* ``stt_model.perception.*`` -> the vendored Conformer.
* everything else (``rnnt_*``, ``tts_model.*``) is skipped lazily.
"""
loaded: set[str] = set()
lm_weights: list[tuple[str, torch.Tensor]] = []
perception_state: dict[str, torch.Tensor] = {}
embed_weight: torch.Tensor | None = None
for name, tensor in weights:
if name.startswith("stt_model.llm."):
lm_weights.append(("backbone." + name.removeprefix("stt_model.llm."), tensor))
elif name == "stt_model.lm_head.weight":
lm_weights.append(("lm_head.weight", tensor))
elif name == "stt_model.embed_tokens.weight":
embed_weight = tensor
elif name == "stt_model.function_head.weight":
with torch.no_grad():
self.function_head.weight.copy_(tensor)
loaded.add("function_head.weight")
elif name.startswith("stt_model.perception."):
perception_state[name.removeprefix("stt_model.perception.")] = tensor
if embed_weight is None:
raise ValueError("Checkpoint is missing 'stt_model.embed_tokens.weight'.")
if not lm_weights:
raise ValueError("Checkpoint is missing the 'stt_model.llm.*' backbone subtree.")
if not perception_state:
raise ValueError("Checkpoint is missing the 'stt_model.perception.*' subtree.")
if "function_head.weight" not in loaded and self._use_function_head:
raise ValueError("Checkpoint is missing 'stt_model.function_head.weight'.")
with torch.no_grad():
self.embed_tokens.weight.copy_(embed_weight)
loaded.add("embed_tokens.weight")
# The vLLM NemotronH model expects an input embedding; the checkpoint's
# shared embed_tokens fills that slot (inputs_embeds bypass it at runtime).
lm_weights.append(("backbone.embeddings.weight", embed_weight))
lm_loaded = self.language_model.load_weights(iter(lm_weights))
loaded |= {f"language_model.{n}" for n in lm_loaded}
missing, unexpected = self.perception.load_state_dict(perception_state, strict=False)
if missing:
raise ValueError(f"NemotronVoiceChat perception is missing weights: {sorted(missing)[:10]} ...")
if unexpected:
logger.warning(
"NemotronVoiceChat perception ignoring %d unexpected tensors (e.g. %s)",
len(unexpected),
sorted(unexpected)[:5],
)
device = self.vllm_config.device_config.device
self.perception.to(device=device, dtype=self._dtype).eval()
self.embed_tokens.to(device=device, dtype=self._dtype)
self.function_head.to(device=device, dtype=self._dtype)
loaded |= {f"perception.{n}" for n in perception_state}
logger.info(
"NemotronVoiceChat thinker loaded: %d backbone tensors, %d perception tensors, "
"fusion weights (text=%.1f, audio=%.1f, function=%.1f)",
len(lm_weights),
len(perception_state),
self._w_text,
self._w_audio,
self._w_func,
)
return loaded
vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py
+85
View File
@@ -0,0 +1,85 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""NemotronVoiceChat pipeline: thinker (speech -> frame-locked text) -> talker
(text timeline -> 31-quantizer RVQ code stacks) -> code2wav (codes -> 22.05 kHz PCM).
``nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`` is an 11B full-duplex speech-to-speech
model. Offline it is served as a qwen3-omni-style 3-stage pipeline:
* Stage 0 (``thinker``, LLM_AR): vendored FastConformer perception + additive
channel fusion feeding the NemotronH backbone on vLLM; emits the frame-locked
agent text channel (also the final text output) and carries the full frame
timeline + metadata to the talker as a latent payload.
* Stage 1 (``talker``, LLM_AR): the vendored Gemma3-based EAR-TTS model stepped
per frame (classifier-free guidance + MoG sampling run inside the vendored
forward); emits one 31-code stack per frame as latents.
* Stage 2 (``code2wav``, LLM_GENERATION): the vendored RVQ-VAE codec decoder
(fp32) producing the final 22.05 kHz waveform.
The STT->TTS coupling is one-directional and token-ids-only (verified against
the NeMo reference), which is what makes the offline cascade exact.
"""
from vllm_omni.config.stage_config import (
PipelineConfig,
StageExecutionType,
StagePipelineConfig,
)
_PROC = "vllm_omni.model_executor.stage_input_processors.nemotron_voicechat"
NEMOTRON_VOICECHAT_PIPELINE = PipelineConfig(
model_type="nemotron_voicechat",
model_arch="NemotronVoiceChatThinkerForConditionalGeneration",
# Named after the alias key so the deploy-yaml stem substring-matches the
# checkpoint directory name (NVIDIA-NemotronLabs-VoiceChat-11B) during
# bare-path auto-detection.
default_deploy_config_name="nemotron_labs_voicechat.yaml",
stages=(
StagePipelineConfig(
stage_id=0,
model_stage="thinker",
execution_type=StageExecutionType.LLM_AR,
input_sources=(),
owns_tokenizer=True,
final_output=True,
final_output_type="text",
engine_output_type="latent",
# thinker -> talker sync mode is token-path only: the frame-locked
# text timeline rides the talker's prompt_token_ids (built by
# thinker2talker_token_only on the talker stage); no connector
# payload and no full-payload wait for stage 1. With
# ``async_chunk: true`` in the deploy yaml, the timeline instead
# streams incrementally over the connector (cumulative ids.all per
# thinker step) so the talker starts before the thinker finishes.
async_chunk_process_next_stage_input_func=f"{_PROC}.thinker2talker_async_chunk",
sampling_constraints={"detokenize": True},
),
StagePipelineConfig(
stage_id=1,
model_stage="talker",
execution_type=StageExecutionType.LLM_AR,
model_arch="NemotronVoiceChatTalkerForConditionalGeneration",
hf_config_name="talker_config",
input_sources=(0,),
engine_output_type="latent",
custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload",
async_chunk_process_next_stage_input_func=f"{_PROC}.talker2code2wav_async_chunk",
sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
sampling_constraints={"detokenize": False},
),
StagePipelineConfig(
stage_id=2,
model_stage="code2wav",
execution_type=StageExecutionType.LLM_GENERATION,
model_arch="NemotronVoiceChatCode2Wav",
hf_config_name="code2wav_config",
input_sources=(1,),
final_output=True,
final_output_type="audio",
engine_output_type="audio",
sync_process_input_func=f"{_PROC}.talker2code2wav_token_only",
sampling_constraints={"detokenize": False},
),
),
)
vllm_omni/model_executor/models/nemotron_voicechat/runtime_info.py
+30
View File
@@ -0,0 +1,30 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Runtime-info helpers shared by the thinker and talker AR stages."""
from __future__ import annotations
from typing import Any
def merge_runtime_info(info_dict: dict[str, Any]) -> dict[str, Any]:
"""Flatten ``additional_information`` into the info dict (top-level wins)."""
additional = info_dict.get("additional_information")
if isinstance(additional, dict):
merged = {k: v for k, v in info_dict.items() if k != "additional_information"}
for k, v in additional.items():
merged.setdefault(k, v)
return merged
return info_dict
def require_request_id(info: dict[str, Any], stage: str) -> str:
"""Resolve the request id from the runtime info (or its meta), or fail."""
request_id = info.get("request_id")
if request_id is None:
meta = info.get("meta")
if isinstance(meta, dict):
request_id = meta.get("request_id")
if request_id is None:
raise RuntimeError(f"NemotronVoiceChat {stage} requires a request_id in the runtime info.")
return str(request_id)
vllm_omni/model_executor/models/registry.py
+17
View File
@@ -176,6 +176,23 @@ _OMNI_MODELS = {
"personaplex_code2wav",
"PersonaPlexCode2Wav",
),
## nemotron_voicechat (NVIDIA-NemotronLabs-VoiceChat-11B, full-duplex S2S;
## 3-stage offline speech->speech pipeline: thinker/talker/code2wav)
"NemotronVoiceChatThinkerForConditionalGeneration": (
"nemotron_voicechat",
"nemotron_voicechat_thinker",
"NemotronVoiceChatThinkerForConditionalGeneration",
),
"NemotronVoiceChatTalkerForConditionalGeneration": (
"nemotron_voicechat",
"nemotron_voicechat_talker",
"NemotronVoiceChatTalkerForConditionalGeneration",
),
"NemotronVoiceChatCode2Wav": (
"nemotron_voicechat",
"nemotron_voicechat_code2wav",
"NemotronVoiceChatCode2Wav",
),
## higgs-audio v2
"HiggsAudioV2ForConditionalGeneration": (
"higgs_audio_v2",
vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py
+355
View File
@@ -0,0 +1,355 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Stage transfer processors for the NemotronVoiceChat 3-stage pipeline.
thinker (0) -> talker (1): token path only. The thinker's frame-locked text
timeline becomes the talker's ``prompt_token_ids``:
``[text_pad_id] * logical_prompt_token_len + generated frame tokens``. The
talker must process the PAD prompt region too (its KV state at the first
acoustic frame depends on it — NeMo trims prompt-region codes only after the
loop), so the prompt slots are reconstructed here, and the timeline metadata
rides ``additional_information``.
talker (1) -> code2wav (2): PersonaPlex-style full payload. The talker
accumulates one 31-quantizer code stack per frame under ``("codes","audio")``;
the producer ships the prompt-trimmed ``[frames, 31]`` stack once the request
finishes, and the sync placeholder builder sizes stage-2's dummy prompt.
"""
from __future__ import annotations
from typing import Any
import torch
from vllm_omni.data_entry_keys import (
CodesStruct,
IdsStruct,
MetaStruct,
OmniPayloadStruct,
)
# Full-sequence tensors are re-shipped whole; the full-payload accumulator must
# replace (not concatenate) them if the producer fires more than once. The AR
# runner flattens nested payloads before accumulation, so the effective key is
# the flattened "codes.audio" (the unflattened root is kept for robustness).
_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset({"codes", "codes.audio", "meta", "meta.nvc_logical_prompt_len"})
# <SPECIAL_12> in the Nemotron-Nano-9B-v2 vocab; the checkpoint config's
# stt pad_token. The thinker also reports it in its latent metadata, which
# takes precedence when present.
_DEFAULT_TEXT_PAD_ID = 12
def _info_get(info: Any, key: str) -> Any:
if isinstance(info, dict):
if key in info:
return info[key]
meta = info.get("meta")
if isinstance(meta, dict):
return meta.get(key)
return None
def thinker2talker_token_only(
source_outputs: list,
prompt: Any = None,
_requires_multimodal_data: bool = False,
) -> list:
"""Build the talker's engine inputs from the thinker's finished outputs.
``prompt_token_ids`` = the full frame-aligned text timeline including the
PAD prompt region. The thinker's vLLM prompt covers
``logical_prompt_token_len + 1`` positions (the +1 is acoustic frame 0), so
``logical_prompt_token_len = len(thinker prompt) - 1`` and the generated
tokens are exactly the acoustic-frame text channel.
"""
from vllm_omni.inputs.data import OmniTokensPrompt
del prompt, _requires_multimodal_data
inputs: list = []
for thinker_output in source_outputs:
if not getattr(thinker_output, "finished", False):
continue
output = thinker_output.outputs[0]
generated = list(getattr(output, "cumulative_token_ids", None) or getattr(output, "token_ids", None) or [])
prompt_ids = list(getattr(thinker_output, "prompt_token_ids", None) or [])
logical_prompt_len = max(len(prompt_ids) - 1, 0)
pad_id = _DEFAULT_TEXT_PAD_ID
info = getattr(thinker_output, "additional_information", None)
reported_pad = _info_get(info, "nvc_text_pad_id")
if reported_pad is not None:
pad_id = int(reported_pad)
timeline = [pad_id] * logical_prompt_len + [int(t) for t in generated]
inputs.append(
OmniTokensPrompt(
# The vLLM-side prompt is one placeholder (timeline position 0;
# the NeMo TTS loop starts at t=1). Each decode step consumes
# one timeline position from additional_information, so the
# talker needs len(timeline) - 1 decode steps.
prompt_token_ids=[0],
additional_information={
"nvc_text_timeline": timeline,
"nvc_logical_prompt_len": logical_prompt_len,
"nvc_frame_count": len(generated),
"nvc_text_pad_id": pad_id,
},
multi_modal_data=None,
mm_processor_kwargs=None,
)
)
return inputs
def thinker2talker_async_chunk(
transfer_manager: Any,
multimodal_output: Any = None,
request: Any = None,
is_finished: bool = False,
**kwargs: Any,
) -> OmniPayloadStruct | None:
"""Streaming producer: ship the CUMULATIVE text timeline each thinker step.
The AR receive path REPLACES the downstream request's
additional_information with each chunk payload, so every chunk must carry
the whole timeline so far: ``[pad] * logical_prompt_len + sampled tokens``
under ``ids.all`` (the talker adopts any longer timeline it sees). The
timeline is a few hundred ints, so cumulative re-ship is cheap and keeps
the payload self-contained. ``meta.num_processed_tokens`` carries the
logical prompt length for the code2wav producer's prompt-region trim.
"""
del multimodal_output, kwargs
request_id = request.external_req_id
prompt_ids = list(getattr(request, "prompt_token_ids", None) or [])
# vLLM prompt = logical prompt + 1 placeholder (acoustic frame 0).
logical_prompt_len = max(len(prompt_ids) - 1, 0)
generated = list(getattr(request, "output_token_ids", None) or [])
pad_id = _DEFAULT_TEXT_PAD_ID
reported_pad = _info_get(getattr(request, "additional_information", None), "nvc_text_pad_id")
if reported_pad is not None:
pad_id = int(reported_pad)
timeline = [pad_id] * logical_prompt_len + [int(t) for t in generated]
# Skip no-progress wakeups (same length as the last emitted chunk).
state = transfer_manager.request_payload.get(request_id)
last_len = int(state.get("nvc_timeline_len", -1)) if isinstance(state, dict) else -1
if len(timeline) <= last_len and not is_finished:
return None
transfer_manager.request_payload[request_id] = {"nvc_timeline_len": len(timeline)}
return OmniPayloadStruct(
ids=IdsStruct(all=timeline),
meta=MetaStruct(
finished=torch.tensor(bool(is_finished), dtype=torch.bool),
num_processed_tokens=logical_prompt_len,
# The talker's vLLM prompt stays a single placeholder token.
next_stage_prompt_len=1,
),
)
def talker2code2wav_async_chunk(
transfer_manager: Any,
multimodal_output: Any = None,
request: Any = None,
is_finished: bool = False,
**kwargs: Any,
) -> OmniPayloadStruct | None:
"""Streaming producer: ship prompt-trimmed CUMULATIVE code stacks in chunks.
Follows the NeMo incremental-decode recipe (``decode_one_audio_step``):
each chunk carries the FULL trimmed code history and
``meta.left_context_size`` = frames the codec already emitted, so the
code2wav stage re-decodes the whole prefix and slices off only the new
samples — no seams, no crossfade, and the concatenated stream equals the
prefix-decode of the final stack. Chunk cadence comes from the connector
config key ``codec_chunk_frames`` (default 13 frames ~= 1 s of audio).
"""
del kwargs
request_id = request.external_req_id
codes = None
if isinstance(multimodal_output, dict):
nested = multimodal_output.get("codes")
codes = nested.get("audio") if isinstance(nested, dict) else None
if codes is None:
codes = multimodal_output.get("codes.audio")
if not isinstance(codes, torch.Tensor) or codes.numel() == 0:
if not is_finished:
return None
return _empty_finished_payload()
if codes.ndim == 1:
codes = codes.reshape(1, -1)
# Prompt-region trim (rows are NeMo timeline steps t=1..; keep t >= P).
# Read the prompt length from the PER-STEP model payload first: the
# request-level additional_information is overwritten both by arriving
# thinker chunks and by talker info updates, so it races. Cache the first
# sighting in the transfer-manager state as a fallback.
state = transfer_manager.request_payload.get(request_id)
prompt_len = _info_get(multimodal_output if isinstance(multimodal_output, dict) else None, "nvc_logical_prompt_len")
if prompt_len is None and isinstance(multimodal_output, dict):
prompt_len = multimodal_output.get("meta.nvc_logical_prompt_len")
if prompt_len is None:
info = getattr(request, "additional_information", None)
prompt_len = _info_get(info, "nvc_logical_prompt_len")
if prompt_len is None:
prompt_len = _info_get(info, "num_processed_tokens")
if prompt_len is None and isinstance(state, dict):
prompt_len = state.get("nvc_prompt_len")
if prompt_len is None:
raise ValueError(
"NemotronVoiceChat talker request is missing its logical prompt length "
"('nvc_logical_prompt_len' or async meta.num_processed_tokens); cannot "
"trim the prompt-region codes for streaming code2wav."
)
trimmed = codes[max(int(prompt_len) - 1, 0) :]
frames_sent = int(state.get("nvc_frames_sent", 0)) if isinstance(state, dict) else 0
new_frames = int(trimmed.shape[0]) - frames_sent
if new_frames <= 0 and not is_finished:
return None
connector = getattr(transfer_manager, "connector", None)
raw_cfg = getattr(connector, "config", {}) or {}
cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
chunk_frames = int(cfg.get("codec_chunk_frames", 13))
if new_frames < chunk_frames and not is_finished:
return None
if new_frames <= 0 and is_finished and frames_sent > 0:
# Everything already shipped; emit a terminal empty chunk so the
# code2wav request observes meta.finished.
return _empty_finished_payload()
transfer_manager.request_payload[request_id] = {
"nvc_frames_sent": int(trimmed.shape[0]),
"nvc_prompt_len": int(prompt_len),
}
return OmniPayloadStruct(
codes=CodesStruct(audio=trimmed.detach().to(dtype=torch.long, device="cpu")),
meta=MetaStruct(
left_context_size=frames_sent,
finished=torch.tensor(bool(is_finished), dtype=torch.bool),
),
)
def _empty_finished_payload() -> OmniPayloadStruct:
return OmniPayloadStruct(
codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
)
def talker2code2wav_full_payload(
transfer_manager: Any = None,
pooling_output: Any = None,
request: Any = None,
is_finished: bool = False,
**kwargs: Any,
) -> OmniPayloadStruct | None:
"""Producer: ship the talker's prompt-trimmed ``[frames, 31]`` code stack.
Codes accumulate per frame under ``("codes","audio")`` in the request's
additional_information; the prompt-region rows were never emitted by the
talker (it only stores codes for generated positions), so no extra trim is
needed here. Ships once, at request finish, with replace semantics.
"""
# NOTE: no is_finished gating — in sync full-payload mode the producer only
# runs at flush time (request already finished), and the flush-time request
# object may still report is_finished()=False. Returning a meta-only
# "pending" struct here would be SENT as the real payload and starve the
# code2wav stage of its codes.
del transfer_manager, is_finished
def _codes_from(src: Any) -> torch.Tensor | None:
if not isinstance(src, dict):
return None
nested = src.get("codes")
audio = nested.get("audio") if isinstance(nested, dict) else None
return audio if audio is not None else src.get("codes.audio")
audio = None
for source in (
getattr(request, "additional_information", None),
getattr(request, "additional_information_cpu", None),
pooling_output,
kwargs.get("multimodal_output"),
):
audio = _codes_from(source)
if audio is not None:
break
if not isinstance(audio, torch.Tensor) or audio.numel() == 0:
import logging
logging.getLogger(__name__).error(
"nemotron_voicechat talker full-payload producer found no codes: "
"pooling_output keys=%s, request.additional_information keys=%s",
sorted(pooling_output.keys()) if isinstance(pooling_output, dict) else type(pooling_output).__name__,
sorted(getattr(request, "additional_information", None) or {})
if isinstance(getattr(request, "additional_information", None), dict)
else type(getattr(request, "additional_information", None)).__name__,
)
return _empty_finished_payload()
if audio.ndim == 1:
audio = audio.reshape(1, -1)
# Rows correspond to NeMo timeline steps t = 1..T-1. Trim the prompt
# region: NeMo keeps gen_codes rows t = P..T-1 (P = logical prompt len),
# i.e. row indices P-1 onward here. P >= 1 always (a system prompt is
# required), so this never drops acoustic frames.
prompt_len = _info_get(getattr(request, "additional_information", None), "nvc_logical_prompt_len")
if prompt_len is None:
prompt_len = _info_get(getattr(request, "additional_information_cpu", None), "nvc_logical_prompt_len")
if prompt_len is None:
# Explicit failure, never a silent trim=0 (untrimmed prompt-region codes
# would ship to the codec). The connector layer logs producer errors.
raise ValueError(
"NemotronVoiceChat talker request is missing 'nvc_logical_prompt_len' in its "
"additional_information; cannot trim the prompt-region codes for code2wav."
)
audio = audio[max(int(prompt_len) - 1, 0) :]
if audio.shape[0] == 0:
return _empty_finished_payload()
# 2-D [frames, Q] payload: the chunk/full-payload transfer adapters route
# >=2-D tensors through the payload channel (1-D would be coerced into
# prompt token ids and dropped from the info merge).
return OmniPayloadStruct(
codes=CodesStruct(audio=audio.detach().to(dtype=torch.long, device="cpu")),
meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
)
def talker2code2wav_token_only(
source_outputs: list,
prompt: Any = None,
_requires_multimodal_data: bool = False,
) -> list:
"""Sync placeholder builder for stage 2 (codes arrive via the connector)."""
from vllm_omni.inputs.data import OmniTokensPrompt
del prompt, _requires_multimodal_data
inputs: list = []
for talker_output in source_outputs:
if not getattr(talker_output, "finished", False):
continue
output = talker_output.outputs[0]
token_ids = getattr(output, "cumulative_token_ids", None) or getattr(output, "token_ids", None) or []
n_frames = len(token_ids)
inputs.append(
OmniTokensPrompt(
prompt_token_ids=[0] * max(n_frames, 1),
additional_information=None,
multi_modal_data=None,
mm_processor_kwargs=None,
)
)
return inputs
__all__ = [
"thinker2talker_token_only",
"thinker2talker_async_chunk",
"talker2code2wav_full_payload",
"talker2code2wav_async_chunk",
"talker2code2wav_token_only",
]
vllm_omni/worker/gpu_ar_model_runner.py
+5
View File
@@ -339,6 +339,11 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
"NemotronHForCausalLM",
"DyninOmniForConditionalGeneration",
"IndexTTS2TalkerForConditionalGeneration",
# nemotron_voicechat: the talker (stage 1) is the full-payload
# producer for code2wav (stage 2); the thinker (stage 0) only
# produces over the connector in async-chunk (streaming) mode.
"NemotronVoiceChatTalkerForConditionalGeneration",
"NemotronVoiceChatThinkerForConditionalGeneration",
}
# The stage-level ``model_arch`` override may be blank so the class
# resolves from the checkpoint's own ``architectures`` (e.g. the Audex
vllm_omni/worker/gpu_generation_model_runner.py
+3
View File
@@ -69,6 +69,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
"CosyVoice3Model",
"DyninOmniForConditionalGeneration",
"IndexTTS2S2MelDecoder",
# nemotron_voicechat: code2wav (stage 2) consumes the talker's
# full-payload code stacks.
"NemotronVoiceChatCode2Wav",
}
if (
getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS