Audio: make FLEURS notebook CUDA-portable

Author: mohammadaaftabv

benchmarking/scripts/audio_fleurs_benchmark.py

@@ -116,6 +116,7 @@ def run_audio_fleurs_benchmark(  # noqa: PLR0913, PLR0915
                 audio_filepath_key="audio_filepath",
                 batch_size=16,
                 fail_on_audio_error=True,
+                adapter_kwargs={"use_cuda_graph_decoder": False},
             ).with_(resources=Resources(gpus=gpus))
         )
         pipeline.add_stage(

nemo_curator/models/asr/nemo_asr.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import gc
+from copy import deepcopy
 from dataclasses import dataclass, field
 from numbers import Integral
 from typing import Any
@@ -71,6 +72,10 @@ class NeMoASRAdapter:
         enable_local_attention: Convert a compatible FastConformer checkpoint
             from global to local attention after loading.
         local_attention_context_size: Left and right local-attention context.
+        use_cuda_graph_decoder: Override NeMo's RNNT CUDA-graph decoder. Leave
+            as ``None`` to preserve the checkpoint default. Set to ``False``
+            on GPU/driver combinations that do not support NeMo's label-loop
+            CUDA graph implementation.
         refresh_cache: Forward NeMo's checkpoint cache refresh flag.
         strict: Forward NeMo's strict checkpoint loading flag.
     """
@@ -80,6 +85,7 @@ class NeMoASRAdapter:
     verbose: bool = False
     enable_local_attention: bool = False
     local_attention_context_size: tuple[int, int] = (128, 128)
+    use_cuda_graph_decoder: bool | None = None
     refresh_cache: bool = False
     strict: bool = True
     _model: Any = field(default=None, init=False, repr=False)
@@ -94,6 +100,9 @@ def __post_init__(self) -> None:
         if not isinstance(self.enable_local_attention, bool):
             msg = "NeMoASRAdapter.enable_local_attention must be a boolean"
             raise TypeError(msg)
+        if self.use_cuda_graph_decoder is not None and not isinstance(self.use_cuda_graph_decoder, bool):
+            msg = "NeMoASRAdapter.use_cuda_graph_decoder must be a boolean or None"
+            raise TypeError(msg)
         try:
             context_size = tuple(self.local_attention_context_size)
         except TypeError as exc:
@@ -132,6 +141,8 @@ def load_model(self, *, num_gpus: int) -> None:
         model = self._load_checkpoint(device)
         if self.enable_local_attention:
             self._configure_local_attention(model)
+        if self.use_cuda_graph_decoder is not None:
+            self._configure_rnnt_cuda_graph_decoder(model)
         self._model = model
 
     def _configure_local_attention(self, model: Any) -> None:  # noqa: ANN401
@@ -154,6 +165,23 @@ def _configure_local_attention(self, model: Any) -> None:  # noqa: ANN401
         )
         change_subsampling_chunking(1)
 
+    def _configure_rnnt_cuda_graph_decoder(self, model: Any) -> None:  # noqa: ANN401
+        """Override the CUDA-graph setting on a compatible NeMo RNNT decoder."""
+        from omegaconf import open_dict
+
+        change_decoding_strategy = getattr(model, "change_decoding_strategy", None)
+        model_cfg = getattr(model, "cfg", None)
+        decoding_cfg = getattr(model_cfg, "decoding", None)
+        greedy_cfg = getattr(decoding_cfg, "greedy", None)
+        if not callable(change_decoding_strategy) or decoding_cfg is None or greedy_cfg is None:
+            msg = f"NeMo checkpoint {self.model_id!r} does not expose a configurable RNNT decoder"
+            raise TypeError(msg)
+
+        decoding_cfg = deepcopy(decoding_cfg)
+        with open_dict(decoding_cfg.greedy):
+            decoding_cfg.greedy.use_cuda_graph_decoder = self.use_cuda_graph_decoder
+        change_decoding_strategy(decoding_cfg=decoding_cfg)
+
     def unload_model(self) -> None:
         """Release worker-local model and CUDA cache state."""
         self._model = None

tests/models/asr/test_nemo_asr.py

@@ -21,6 +21,7 @@
 
 import numpy as np
 import pytest
+from omegaconf import OmegaConf
 
 from nemo_curator.models.asr import nemo_asr
 from nemo_curator.models.asr.base import ASRAdapter
@@ -95,6 +96,25 @@ def test_load_model_configures_local_attention_when_enabled() -> None:
     model.change_subsampling_conv_chunking_factor.assert_called_once_with(1)
 
 
+@pytest.mark.parametrize("enabled", [False, True])
+def test_load_model_configures_rnnt_cuda_graph_decoder_when_requested(enabled: bool) -> None:
+    adapter = NeMoASRAdapter(use_cuda_graph_decoder=enabled)
+    model = _mock_model([])
+    model.cfg = OmegaConf.create({"decoding": {"strategy": "greedy_batch", "greedy": {}}})
+
+    with patch.object(adapter, "_load_checkpoint", return_value=model):
+        adapter.load_model(num_gpus=0)
+
+    decoding_cfg = model.change_decoding_strategy.call_args.kwargs["decoding_cfg"]
+    assert decoding_cfg.strategy == "greedy_batch"
+    assert decoding_cfg.greedy.use_cuda_graph_decoder is enabled
+
+
+def test_nemo_adapter_rejects_invalid_cuda_graph_decoder_value() -> None:
+    with pytest.raises(TypeError, match="use_cuda_graph_decoder must be a boolean or None"):
+        NeMoASRAdapter(use_cuda_graph_decoder="false")  # type: ignore[arg-type]
+
+
 def test_transcribe_batch_uses_one_exact_nemo_batch() -> None:
     model = _mock_model([SimpleNamespace(text="alpha"), SimpleNamespace(text="beta")])
     adapter = NeMoASRAdapter(num_workers=2)

tutorials/audio/fleurs/README.md

@@ -81,6 +81,7 @@ python tutorials/audio/fleurs/main.py \
 | `data_split` | FLEURS split: `train`, `dev`, or `test` |
 | `wer_threshold` | Keep samples with `wer_pct ≤` this value (default: `5.5`) |
 | `stages.1.model_id` | NeMo ASR model for inference |
+| `stages.1.adapter_kwargs.use_cuda_graph_decoder` | RNNT CUDA-graph decoder override. The FLEURS hybrid-RNNT default is `false` for broad driver compatibility. |
 | `stages.1.resources.gpus` | GPUs for ASR (`0` for CPU) |
 | `backend` | `xenna` (default) or `ray_data` |
 
@@ -233,6 +234,7 @@ finally:
 |---|---|---|
 | Output directory already exists | Previous run left `${raw_data_dir}/result/${lang}/` | Remove the directory before re-running |
 | OOM during ASR inference | GPU VRAM too small for model + batch | Reduce `stages.1.batch_size` or use a smaller model |
+| `CUDA error: invalid argument` in RNNT label-loop decoding | NeMo CUDA-graph decoder is unsupported by the local CUDA runtime/driver combination | Set `stages.1.adapter_kwargs.use_cuda_graph_decoder=false` (the supplied FLEURS config already does this) |
 | CPU inference very slow | CPU is 10–50x slower than GPU | Set `stages.1.resources.gpus=1`; CPU is only for testing |
 | Empty output JSONL | `wer_threshold` too strict for the model+language pair | Increase `wer_threshold` or use a better-matching ASR model |
 | HuggingFace download fails | Network/auth issue | Check connectivity; some splits may need `huggingface-cli login` |