nvidia/Nemotron-Labs-Audex-2B

Nemotron-Labs-Audex-2B

Introduction

We're excited to introduce Nemotron-Labs-Audex-2B, a unified audio-text LLM with a similar recipe as Nemotron-Labs-Audex-30B-A3B. Audex-2B extends the vocabulary for discrete audio tokens used for speech and general audio outputs, as well as an audio encoder for speech and general audio inputs. Audex-2B delivers strong abilities on audio tasks (audio understanding, speech recognition and translation, text-to-speech, audio generation, and speech-to-speech generation) while preserving very compelling reasoning, alignment, knowledge, long-context, and agentic capabilities of its text-only LLM backbone with marginal or no regression. Audex-2B operates in both thinking and instruct (non-thinking) modes.

Model Architecture

Templates

Multi-Stage-SFT and Cascaded-RL Pipelines

Quick Start

• Audex-2B follows the ChatML template and supports both thinking and instruct (non-thinking) modes. Reasoning content is enclosed within ` and tags. To activate the instruct (non-thinking) mode, we prepend ` to the beginning of the assistant’s response.

• Audex-2B supports up to a 128K-token context length.

• Audex-2B follows Nemotron-Cascade-2 on text evaluation.

• Audex-2B has different recommended inference setups per audio-related task as described below.

Environment

We use vLLM 0.20.0 container image: vllm/vllm-openai:v0.20.0-cu129

• vLLM inference — text-only reasoning, text-to-speech, text-to-audio, and audio understanding / speech recognition / speech translation: runs on vLLM 0.20.0.
• Hugging Face / transformers inference — requires transformers >= 4.53.0 (tested with 4.53.3) and also works with transformers >= 5.0.

Audio extras: ``vllm/vllm-openai:v0.20.0 image does not include audio codecs. This command installs audio-related packages: python3 -m pip install "vllm[audio]"``.

vLLM plugin: Audex-2B is served through a small vLLM plugin. From the model root, register it once (needed for every vLLM task below): ``pip install -e nemotron_dense_vllm_plugin --no-deps --no-build-isolation``.

Audio QA Inference

Audio QA includes audio understanding, speech recognition, and speech translation (see templates in Introduction).

• vLLM (recommended) — offline ``LLM.generate and an OpenAI-compatible audio_url`` server.
• Hugging Face / transformers — requires transformers >= 4.53.0 (we tested with 4.53.3), and also works with transformers >= 5.

Inputs

To prepare inputs, create a JSON file in the following format with the ``\n`` placeholder:

[
{
"id": "sample_0",
"sound": "/path/to/audio_0.wav",
"conversations": [
{"from": "human", "value": "\nDescribe the audio in detail."},
{"from": "gpt", "value": "N/A"}
]
},
{
"id": "sample_1",
"sound": "/path/to/audio_1.wav",
"conversations": [
{"from": "human", "value": "\n{prompt}"},
{"from": "gpt", "value": "N/A"}
]
},
...
]

Inference recipes

• For audio understanding, we use top_p=0.9 and temperature=0.7.
• For speech recognition and translation, we use greedy sampling.

vLLM (recommended)

To install environments:

python3 -m pip install "vllm[audio]" # audio input decoding; skip if your image already bundles it (see Environment)
pip install -e inference_scripts_vllm/audioqa_scripts --no-deps --no-build-isolation

• Offline (JSON in → JSONL out):

python inference_scripts_vllm/audioqa_scripts/run_audioqa_vllm.py \
--model-path "$(pwd)/checkpoint_folder_full" \
--input-json ./inputs.json \
--output-jsonl ./audioqa_outputs/results.jsonl \
--tensor-parallel-size 1

• OpenAI-compatible server + client:

bash inference_scripts_vllm/audioqa_scripts/serve_audioqa_vllm.sh "$(pwd)/checkpoint_folder_full" 8000
python inference_scripts_vllm/audioqa_scripts/client_audioqa.py --audio /path/to/audio.wav --prompt "Describe this audio."

Hugging Face / transformers

• Example inference script: ``bash inference_scripts_hf/inference_example.sh``.
• Task instruction examples: audio understanding — a question about the audio; speech recognition — ``Transcribe the speech in the input audio.\n; speech translation — a translation instruction such as Translate the speech in the input audio into English.\n``.

Audio Generation Inference

Audio generation includes text-to-speech and text-to-audio generation.

First, prepare vLLM inference using ``bash model_conversion_scripts/prepare_audiogen_vllm_checkpoint.sh (which only creates symlinks of safetensors under checkpoint_folder_audiogen``).

Text-to-audio (TTA)

Download XCodec1 (hf-audio/xcodec-hubert-general-balanced) via

hf download hf-audio/xcodec-hubert-general-balanced --local-dir /path/to/xcodec1

Prepare a folder ``/path/to/caption_txt_dir/ with all .txt files where each contains one caption. Run cd inference_scripts_vllm/audiogen_scripts/ and set --tensor-parallel-size`` to the number of GPUs. Run

XCODEC1_PATH=/path/to/xcodec1 python3 run_audio_gen_vllm_rvq_logit_mask.py \
--task tta \
--model-path $(pwd)/../../checkpoint_folder_audiogen/ \
--dataset-path /path/to/caption_txt_dir/ \
--output-dir ../../tta_outputs/dataset_name/ \
--tensor-parallel-size 1 \
--temperature 1.0 \
--top-k 80 \
--max-tokens 2048 \
--cfg-scale 3.0 \
--cfg-pairs-per-batch 2

Finally (optional), apply the 48 kHz enhancement VAE to the generated waveforms; see ``enhancement_VAE/README.md``.

Text-to-speech (TTS)

we recommend using the standalone Audex causal speech decoder in ``audex_causal_speech_decoder`` (default).

./run_tts_vllm.sh --transcription "The weather is so good, and I want to enjoy the beautiful morning in the park." \
--output-dir ./tts_outputs --utt-id the_weather_is_so_good

Alternatively, users can download the original XCodec2 via this repo and decode the tokens after full generation; this has better quality but is not streaming.

Text-Only Reasoning

The text reasoning follows Nemotron-Cascade-2-30B-A3B.

• To create the...