File size: 11,043 Bytes
ea4bf5c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
---
library_name: transformers
license: other
license_name: nvidia-oneway-noncommercial-license
license_link: https://huggingface.co/nvidia/Nemotron-Labs-Audex-2B/blob/main/license/NVIDIA-OneWay-Noncommercial-License.docx
pipeline_tag: text-generation
track_downloads: true
language:
  - en
tags:
  - nvidia
  - nemotron-labs-audex
  - reasoning
  - general-purpose
  - SFT
  - audio-language-modeling
  - audio-understanding
  - text-to-speech
  - text-to-audio
  - speech-recognition
  - speech-translation
---

# Nemotron-Labs-Audex-2B

<p align="center">

[![Technical Report](https://img.shields.io/badge/2607.05196-Technical_Report-blue)](https://arxiv.org/abs/2607.05196)
[![Models](https://img.shields.io/badge/🤗-Models-blue)](https://huggingface.co/collections/nvidia/nemotron-labs-audex)
</p>

<img src="assets/main.png" alt="main_fig" style="width: 1000px; max-width: 75%;" />


## Introduction
We're excited to introduce [Nemotron-Labs-Audex-2B](https://huggingface.co/nvidia/Nemotron-Labs-Audex-2B), a unified audio-text LLM with a similar recipe as [Nemotron-Labs-Audex-30B-A3B](https://huggingface.co/nvidia/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

<img src="assets/architecture.png" alt="architecture" style="width: 1000px; max-width: 75%;" />

### Templates

<img src="assets/template.png" alt="template" style="width: 1000px; max-width: 75%;" />


### Multi-Stage-SFT and Cascaded-RL Pipelines

<img src="assets/training.png" alt="training" style="width: 1000px; max-width: 75%;" />


### Audex-2B Details

Audex-2B is a compact, dense model trained using the same multi-stage SFT recipe as Audex-30B-A3B. Its smaller size reduces inference resource requirements. Audex-2B is released after multi-stage SFT, while Audex-30B-A3B additionally undergoes cascaded RL.


## Quick Start

- Audex-2B follows the ChatML template and supports both thinking and instruct (non-thinking) modes. Reasoning content is enclosed within `<think>` and `</think>` tags. To activate the instruct (non-thinking) mode, we prepend `<think></think>` 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](https://hub.docker.com/layers/vllm/vllm-openai/v0.20.0-cu129/images/sha256-f4ace3494896eeda800dee284d1fc42ca7f5626f31ceae8e24d1383d770567c2)

- **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 ```<sound>\n``` placeholder:
```json
[
  {
    "id": "sample_0",
    "sound": "/path/to/audio_0.wav",
    "conversations": [
      {"from": "human", "value": "<sound>\nDescribe the audio in detail."},
      {"from": "gpt", "value": "N/A"}
    ]
  },
  {
    "id": "sample_1",
    "sound": "/path/to/audio_1.wav",
    "conversations": [
      {"from": "human", "value": "<sound>\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:
```bash
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):
```bash
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
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<sound>```; speech translation — a translation instruction such as ```Translate the speech in the input audio into English.\n<sound>```.



### 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](https://huggingface.co/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](https://github.com/zhenye234/X-Codec-2.0) 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](https://huggingface.co/nvidia/Nemotron-Cascade-2-30B-A3B). 

- To create the checkpoint that completely matches their setup, run ```python model_conversion_scripts/convert_full_HF_to_textonly_HF.py``` to remove the audio-related vocabularies. 
- Then, the text-only inference completely follows Nemotron-Cascade-2-30B-A3B. See a simple reasoning example with ```cd inference_scripts_vllm/textonly_scripts/; python run_text_vllm_example.py --model-path $(pwd)/../../checkpoint_folder_textonly```.
- Note: you could also avoid model conversion by using ```sampling_params = SamplingParams(allowed_token_ids=list(range(131072)))``` in vLLM inference to mask the audio tokens, although we did not thoroughly test this approach.

### Demo: speech-to-speech interaction with text reasoning
See `inference_scripts_vllm/unified_s2s_scripts/README.md`.

## Reproducibility

The benchmark numbers below use the following setups:

- **Text-to-speech (TTS):** the original non-streaming XCodec2 decoder.
- **Text-to-audio (TTA):** XCodec1 followed by the enhancement VAE.
- **Text:** same as Nemotron-Cascade-2.
- **Audio understanding:** transformers 4.53.3 and Megatron-LM's native inference. 

## Detailed Benchmark Results

### Text Results
<img src="assets/result_text.png" alt="result_text" style="width: 1000px; max-width: 75%;" />
<br></br>


### Text-To-Speech Results
<img src="assets/result_tts.png" alt="result_tts" style="width: 1000px; max-width: 75%;" />
<br></br>


### Text-To-Audio Results
<img src="assets/result_tta.png" alt="result_tta" style="width: 1000px; max-width: 75%;" />
<br></br>


### Speech Recognition and Translation Results
<img src="assets/result_openasr.png" alt="result_openasr" style="width: 1000px; max-width: 75%;" />
<br></br>
<img src="assets/result_noisyasr.png" alt="result_noisyasr" style="width: 1000px; max-width: 75%;" />
<br></br>
<img src="assets/result_multilingual.png" alt="result_multilingual" style="width: 1000px; max-width: 75%;" />
<br></br>


### Audio Understanding Results
<img src="assets/result_audioqa.png" alt="result_audioqa" style="width: 1000px; max-width: 75%;" />
<br></br>

### Speech-To-Speech Results
<img src="assets/result_s2s.png" alt="result_s2s" style="width: 1000px; max-width: 75%;" />
<br></br>


## Release Date
June 8, 2026


## License
Your use of this model is governed by the [NVIDIA Oneway Noncommercial License](license/NVIDIA-OneWay-Noncommercial-License.docx)


## Citation
```
@article{Nemotron-Labs-Audex,
  title={Unified Audio Intelligence Without Regressing on Text Intelligence},
  author={Kong, Zhifeng and Lee, Sang-gil and Kim, Jaehyeon and Wang, Boxin and Liu, Zihan and Kim, Sungwon and Chen, Yang and Goel, Arushi and Roy, Rajarshi and Dai, Wenliang and Yang, Zhuolin and Chen, Yangyi and Jiang, Dongfu and Ghosh, Sreyan and Rintamaki, Tuomas and Tao, Andrew and Raiman, Jonathan and Shoeybi, Mohammad and Catanzaro, Bryan and Ping, Wei},
  year={2026}
}
```