Skip to content

Commit a33342d

Browse files
mergennachinmergennachin
authored
[examples] Add Nemotron 3 Diarization with MLX (pytorch#23137)
Export the Transformers model with BF16 by default and run offline or streaming diarization with a native C++ runner. Keep graph capture and streaming logic independent of MLX-specific build and lowering. pytorch#23131 Test Plan: - Exported BF16 with full MLX delegation; checked captured graphs against Transformers eager execution across dynamic input lengths. - Ran native ExecuTorch audio smoke tests across all four streaming presets, including reset, feed partitioning, final padding, and invalid inputs. Compared with previous ExecuTorch outputs: maximum probability difference 7.63e−6, with zero speaker-decision changes. - Passed local UBSan runner regression cases, plus invalid-metadata and execution-failure recovery checks. Authored with OpenAI Codex.
1 parent 42b7709 commit a33342d

10 files changed

Lines changed: 998 additions & 0 deletions

File tree

‎.github/workflows/mlx.yml‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ on:
1717
- extension/audio/**
1818
- examples/models/gemma4_31b/**
1919
- examples/models/muse-glimmer/**
20+
- examples/models/nemotron3_diarization/**
2021
- examples/models/parakeet/**
2122
- examples/models/supertonic/**
2223
- examples/models/voxtral_realtime/**
@@ -354,6 +355,31 @@ jobs:
354355
echo "Success: built solo_runner, dflash_runner and muse_glimmer_worker"
355356
echo "::endgroup::"
356357
358+
test-mlx-nemotron3-diarization-build:
359+
needs: run-decision
360+
if: |
361+
github.event_name == 'pull_request' ||
362+
needs.run-decision.outputs.is-full-run == 'true'
363+
uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
364+
with:
365+
default-packages: ""
366+
job-name: test-mlx-nemotron3-diarization-build
367+
runner: macos-14-xlarge
368+
python-version: "3.12"
369+
submodules: recursive
370+
ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
371+
timeout: 90
372+
script: |
373+
set -eux
374+
${CONDA_RUN} python install_executorch.py > /dev/null
375+
${CONDA_RUN} make nemotron3-diarization-mlx
376+
BIN_DIR=cmake-out/examples/models/nemotron3_diarization
377+
test -x "${BIN_DIR}/nemotron3_diarization_runner"
378+
test -f "${BIN_DIR}/mlx.metallib"
379+
# Ensure the shared install still supplies the speech/LLM runner targets.
380+
test -f cmake-out/lib/libextension_asr_runner.a
381+
test -f cmake-out/lib/libextension_llm_runner.a
382+
357383
backend-tester:
358384
needs: run-decision
359385
if: |

‎Makefile‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
# - whisper: Speech recognition model (CPU, CUDA, Metal)
2020
# - parakeet: Speech recognition model (CPU, CUDA, Metal, MLX)
2121
# - sortformer: Speaker diarization model (CPU, CUDA)
22+
# - nemotron3-diarization: Speaker diarization model (MLX)
2223
# - supertonic: Text-to-speech model (MLX)
2324
# - silero_vad: Voice activity detection model (CPU)
2425
# - llama: Text generation model (CPU)
@@ -121,6 +122,7 @@ help:
121122
@echo " dinov2-cuda-debug - Build DINOv2 runner with CUDA backend (debug mode)"
122123
@echo " sortformer-cuda - Build Sortformer runner with CUDA backend"
123124
@echo " sortformer-cpu - Build Sortformer runner with CPU backend"
125+
@echo " nemotron3-diarization-mlx - Build Nemotron 3 Diarization runner with MLX backend"
124126
@echo " supertonic-mlx - Build Supertonic runner with MLX backend"
125127
@echo " silero-vad-cpu - Build Silero VAD runner with CPU backend"
126128
@echo " llama-cuda - Build Llama runner with CUDA backend"
@@ -285,6 +287,12 @@ dinov2-cuda-debug:
285287
@echo "✓ Build complete!"
286288
@echo " Binary: cmake-out/examples/models/dinov2/dinov2_runner"
287289

290+
.PHONY: nemotron3-diarization-mlx
291+
nemotron3-diarization-mlx:
292+
cmake --preset mlx-release -DEXECUTORCH_BUILD_XNNPACK=ON
293+
cmake --build --preset mlx-release-install
294+
cd examples/models/nemotron3_diarization && cmake --workflow --preset nemotron3-diarization-mlx
295+
288296
sortformer-cuda:
289297
@echo "==> Building and installing ExecuTorch with CUDA..."
290298
cmake --workflow --preset llm-release-cuda

‎examples/models/nemotron3_diarization/CMakeLists.txt‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
cmake_minimum_required(VERSION 3.24)
8+
project(nemotron3_diarization_runner)
9+
10+
set(CMAKE_CXX_STANDARD 17)
11+
set(CMAKE_CXX_STANDARD_REQUIRED ON)
12+
set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..)
13+
include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake)
14+
15+
set(gflags_DIR ${CMAKE_CURRENT_BINARY_DIR}/../../../third-party/gflags)
16+
find_package(gflags REQUIRED)
17+
find_package(executorch CONFIG REQUIRED FIND_ROOT_PATH_BOTH)
18+
if(NOT TARGET mlxdelegate
19+
OR NOT TARGET mlx
20+
OR NOT TARGET xnnpack_backend
21+
)
22+
message(
23+
FATAL_ERROR
24+
"Build and install ExecuTorch with EXECUTORCH_BUILD_MLX=ON and EXECUTORCH_BUILD_XNNPACK=ON"
25+
)
26+
endif()
27+
executorch_target_link_options_shared_lib(mlxdelegate)
28+
executorch_target_link_options_shared_lib(xnnpack_backend)
29+
30+
add_library(nemotron3_diarization STATIC runner.cpp)
31+
target_include_directories(nemotron3_diarization PUBLIC ${EXECUTORCH_ROOT}/..)
32+
target_link_libraries(
33+
nemotron3_diarization
34+
PUBLIC executorch
35+
extension_module
36+
extension_data_loader
37+
extension_tensor
38+
extension_flat_tensor
39+
optimized_native_cpu_ops_lib
40+
cpublas
41+
eigen_blas
42+
xnnpack_backend
43+
XNNPACK
44+
xnnpack-microkernels-prod
45+
)
46+
if(TARGET kleidiai)
47+
target_link_libraries(nemotron3_diarization PUBLIC kleidiai)
48+
endif()
49+
executorch_target_link_options_shared_lib(optimized_native_cpu_ops_lib)
50+
add_executable(nemotron3_diarization_runner main.cpp)
51+
target_include_directories(
52+
nemotron3_diarization_runner
53+
PRIVATE ${EXECUTORCH_ROOT}/third-party/json/include
54+
)
55+
target_link_libraries(
56+
nemotron3_diarization_runner PRIVATE nemotron3_diarization gflags mlxdelegate
57+
mlx
58+
)
59+
executorch_target_copy_mlx_metallib(nemotron3_diarization_runner)

‎examples/models/nemotron3_diarization/CMakePresets.json‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
{
2+
"version": 6,
3+
"configurePresets": [
4+
{
5+
"name": "nemotron3-diarization-mlx",
6+
"displayName": "Nemotron 3 Diarization runner (MLX)",
7+
"binaryDir": "${sourceDir}/../../../cmake-out/examples/models/nemotron3_diarization",
8+
"cacheVariables": {
9+
"EXECUTORCH_BUILD_MLX": "ON",
10+
"CMAKE_BUILD_TYPE": "Release",
11+
"CMAKE_OSX_DEPLOYMENT_TARGET": "14.0",
12+
"CMAKE_FIND_ROOT_PATH": "${sourceDir}/../../../cmake-out",
13+
"CMAKE_PREFIX_PATH": "${sourceDir}/../../../cmake-out"
14+
},
15+
"condition": {"type": "equals", "lhs": "${hostSystemName}", "rhs": "Darwin"}
16+
}
17+
],
18+
"buildPresets": [
19+
{
20+
"name": "nemotron3-diarization-mlx",
21+
"configurePreset": "nemotron3-diarization-mlx",
22+
"configuration": "Release"
23+
}
24+
],
25+
"workflowPresets": [
26+
{
27+
"name": "nemotron3-diarization-mlx",
28+
"steps": [
29+
{"type": "configure", "name": "nemotron3-diarization-mlx"},
30+
{"type": "build", "name": "nemotron3-diarization-mlx"}
31+
]
32+
}
33+
]
34+
}

‎examples/models/nemotron3_diarization/README.md‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
# Nemotron 3 Diarization on MLX
2+
3+
```bash
4+
# From the repository root on Apple Silicon:
5+
python -m pip install git+https://github.com/huggingface/transformers "librosa>=0.10"
6+
7+
# Export (approximate BF16; add --dtype fp32 for comparison with Transformers)
8+
# Uses the Transformers frontend.
9+
python -m executorch.examples.models.nemotron3_diarization.export_nemotron \
10+
--hf-model nvidia/Nemotron-3-Diarization \
11+
--output-dir nemotron_exports
12+
13+
# Build
14+
make nemotron3-diarization-mlx
15+
16+
# Convert to mono 16 kHz PCM16 and run
17+
ffmpeg -i audio.wav -ac 1 -ar 16000 -c:a pcm_s16le audio_16khz_mono.wav
18+
cmake-out/examples/models/nemotron3_diarization/nemotron3_diarization_runner \
19+
--model_path=nemotron_exports/nemotron3_diarization.pte \
20+
--audio_path=audio_16khz_mono.wav \
21+
--preset=offline \
22+
--output=segments.json
23+
```

‎examples/models/nemotron3_diarization/__init__.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.

0 commit comments

Comments
 (0)