NeMo 2.6 incompatible with `numba-cuda`

[ChrisRahme-Innovum]

Describe the bug

NeMo 2.6.0+ has an incompatible dependency specification for numba-cuda. The package requires numba-cuda==0.15.1, but the code attempts to import redirect_numba_module from numba.cuda.utils, which doesn't exist in version 0.15.1. This causes an ImportError when trying to import NeMo ASR models.

Steps/Code to reproduce bug

1. Install NeMo 2.6 with ASR on GPU:

pip install "nemo-toolkit[asr]==2.6.1"

2. Verify numba-cuda version:

python -m pip list | grep numba-cuda
# numba-cuda 0.15.1

3. Attempt to import any ASR model:

from nemo.collections.asr.models.msdd_models import NeuralDiarizer

4. Observe the error:

ImportError: cannot import name 'redirect_numba_module' from 'numba.cuda.utils'

Expected behavior

NeMo should install with compatible dependencies and allow importing ASR models without errors. The import should succeed, or the dependency specification should require a compatible version of numba-cuda.

Environment overview (please complete the following information)

• Environment location: Docker (nvidia/cuda:12.2.2-cudnn8-runtime-ubuntu22.04 base image)
• Method of NeMo install: pip install via Poetry
• Install command: poetry add "nemo-toolkit[asr]"

Environment details

• OS version: Ubuntu 22.04
• PyTorch version: 2.8.0
• Python version: 3.12.11
• CUDA version: 12.2.2
• cuDNN version: 8

Additional context

We were using NeMo 2.4.0. Upgrading to 2.6 also allowed Numpy to be upgraded from 1.x to 2.x, which may be responsible for the issue, since Poetry logged the following:

nemo-toolkit[asr] (>=2.6.0,<3.0.0) requires numba-cuda (0.15.1) or numpy (>=1.22,<2.0.0)

With numpy >= 2.0.0, this meant that numba-cuda has to be at version 0.15.1, which does not export numba.cuda.utils.redirect_numba_module.

Full traceback:

  File "diarize.py", line 17, in <module>
    from nemo.collections.asr.models.msdd_models import NeuralDiarizer
  File "nemo/collections/asr/__init__.py", line 15, in <module>
    from nemo.collections.asr import data, losses, models, modules
  File "nemo/collections/asr/models/__init__.py", line 15, in <module>
    from nemo.collections.asr.models.aed_multitask_models import EncDecMultiTaskModel
  File "nemo/collections/asr/models/aed_multitask_models.py", line 32, in <module>
    from nemo.collections.asr.metrics import MultiTaskMetric
  File "nemo/collections/asr/metrics/__init__.py", line 15, in <module>
    from nemo.collections.asr.metrics.bleu import BLEU
  File "nemo/collections/asr/metrics/bleu.py", line 24, in <module>
    from nemo.collections.asr.parts.submodules.multitask_decoding import AbstractMultiTaskDecoding
  File "nemo/collections/asr/parts/submodules/multitask_decoding.py", line 22, in <module>
    from nemo.collections.asr.parts.submodules.multitask_beam_decoding import (
  File "nemo/collections/asr/parts/submodules/multitask_beam_decoding.py", line 22, in <module>
    from nemo.collections.asr.modules.transformer import (
  File "nemo/collections/asr/modules/__init__.py", line 15, in <module>
    from nemo.collections.asr.modules.audio_preprocessing import (
  File "nemo/collections/asr/modules/audio_preprocessing.py", line 24, in <module>
    from nemo.collections.asr.parts.numba.spec_augment import SpecAugmentNumba, spec_augment_launch_heuristics
  File "nemo/collections/asr/parts/numba/__init__.py", line 15, in <module>
    from nemo.collections.asr.parts.numba.rnnt_loss.rnnt_pytorch import RNNTLossNumba
  File "nemo/collections/asr/parts/numba/rnnt_loss/__init__.py", line 15, in <module>
    from nemo.collections.asr.parts.numba.rnnt_loss.rnnt import rnnt_loss_cpu, rnnt_loss_gpu
  File "nemo/collections/asr/parts/numba/rnnt_loss/rnnt.py", line 32, in <module>
    from numba import cuda
  File "numba_cuda/numba/cuda/__init__.py", line 49, in <module>
    from .device_init import *
  File "numba_cuda/numba/cuda/device_init.py", line 3, in <module>
    from numba.cuda import cg
  File "numba_cuda/numba/cuda/cg.py", line 5, in <module>
    from numba.cuda.extending import intrinsic
  File "numba_cuda/numba/cuda/extending.py", line 6, in <module>
    from numba.cuda.models import register_model  # noqa: F401
  File "numba_cuda/numba/cuda/models.py", line 8, in <module>
    from numba.cuda.types import Dim3, GridGroup, CUDADispatcher
  File "numba_cuda/numba/cuda/types/__init__.py", line 9, in <module>
    from .abstract import *
  File "numba_cuda/numba/cuda/types/abstract.py", line 5, in <module>
    from numba.cuda.utils import redirect_numba_module
ImportError: cannot import name 'redirect_numba_module' from 'numba.cuda.utils' (.venv/lib/python3.12/site-packages/numba_cuda/numba/cuda/utils.py)

[pzelasko]

@nithinraok I see you pinned the numba-cuda version to 0.15.1 in #15166 even though it was <=0.20.0 before, and the latest version today is 0.26.0. What was your motivation for that?

In the meantime, @ChrisRahme-Innovum, can you try to upgrade pip install --upgrade numba-cuda to a newer version? While we figure it out on our end.

[nithinraok]

Yes, we faced similar issue last timd and had to pin to specific version. <=0.20 was installing 0.20 which was raising issues.

[pzelasko]

I'm afraid we'll have to defer addressing it for now - and return to this when we are revising and trimming NeMo dependency list. This is planned soon.