Fix `numba-cuda` and `cuda-python` installation and usage by artbataev · Pull Request #15506 · NVIDIA-NeMo/NeMo

artbataev · 2026-03-17T11:11:40Z

Important

The Update branch button must only be pressed in very rare occassions.
An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

Collection: [Note which collection this PR will affect]

Changelog

Add specific line by line info of high level changes in this PR.

Usage

You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?
Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

New Feature
Bugfix
Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

Related to # (issue)

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

pzelasko · 2026-03-17T13:48:11Z

requirements/requirements.txt

 huggingface_hub>=0.24
 numba ; platform_system == 'Darwin'
-numba-cuda==0.15.1 ; platform_system != 'Darwin'
+numba-cuda[cu13] ; platform_system != 'Darwin'


Thinking from first principles:

we have CU13 users that have cutting edge hardware / want cutting edge perf

we have CU12 users that don't have cutting edge hardware

we likely have CU11 users too

we have CPU-only env users, also on platforms like MacOS

How do we cater to all of them? Sane default:

pip install nemo-toolkit - does not pull any CUDA depedencies; the codebase treats all of them as optional and raises an informative error about what to install if you try to use them

CUDA version is a bit more problematic because it depends on what pytorch was built against. We can't really control this in a robust way. So let's give options. Is it possible to do sth like this?

pip install nemo-toolkit[cu13] - pulls CUDA 13 deps

pip install nemo-toolkit[cu12] - pulls CUDA 12 deps

pip install nemo-toolkit[cu11] - pulls CUDA 11 deps

FWIW Curator also currently only supports CUDA 12, and Curator depends on this. So +1-ing that recommendation, Curator can then depend on nemo-toolkit[cu12]

Regarding CUDA 11.x - I think it is ok if the users of CUDA 11.x will install extra dependencies manually.
I'm unsure about the best numba version for this case, but latest versions do not support 11.x.
For cuda-python, we require at least 12.3.x, better to use 12.6.x. So, for CUDA 11.x it is useless.

Since both libraries are optional, let's for now keep only cu12 and cu13 extras.

Since both libraries are optional, let's for now keep only cu12 and cu13 extras.

OK

nithinraok

can we add cuda-python to the reqs
(nemo/core/utils/cuda_python_utils.py)

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Signed-off-by: artbataev <artbataev@users.noreply.github.com>

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

pzelasko · 2026-03-20T15:15:48Z

nemo/core/utils/optional_libs.py

 """

-__all__ = ["KENLM_AVAILABLE", "K2_AVAILABLE", "TRITON_AVAILABLE", "kenlm_required", "k2_required", "triton_required"]
+__all__ = [


Now I'm wondering if we should have API like this instead:

k2 = ext_k2.maybe_import() # returns python module or None @ext_k2.required() def foo(): ... if ext_k2.available(): ...

possibly implemented like this:

class OptionalDependency: NAME = "" INSTALATION_MESSAGE = "" def required(self): ... # moved from _lib_required def available(self): ... # moved from is_lib_available class ext_k2(OptionalDependency): NAME = "k2" INSTALLATION_MESSAGE = "..."

k2 = ext_k2.maybe_import()

I initially thought about this, but this approach imposes makes imports more strict. You cannot anymore use from k2 import ... etc

class OptionalDependency: NAME = "" INSTALATION_MESSAGE = "" def required(self): ... # moved from _lib_required def available(self): ... # moved from is_lib_available class ext_k2(OptionalDependency): NAME = "k2" INSTALLATION_MESSAGE = "..."

This approach has more overhead for required due to dynamic availability check. I prefer to check once -> create a constant.

OK, maybe I'm overthinking this :)

pzelasko · 2026-03-20T15:20:46Z

requirements/requirements.txt

 huggingface_hub>=0.24
 numba ; platform_system == 'Darwin'
-numba-cuda[cu13] ; platform_system != 'Darwin'
+numba-cuda ; platform_system != 'Darwin'


Should this be removed from main requirements? Doesn't it conflict with lines added by cu12/cu13?

pzelasko · 2026-03-20T15:21:05Z

setup.py

 extras_require['all'] = list(chain(val for key, val in extras_require.items()))

+# CUDA version extras (not included in 'all' - user must explicitly select)
+extras_require['cu12'] = [


For consistency, can we move to requirements_cu12.txt (same with cu13)?

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

pzelasko · 2026-03-20T16:05:36Z

nemo/collections/asr/modules/audio_preprocessing.py

 from nemo.utils import logging, logging_mode

+if NUMBA_CUDA_AVAILABLE:
+    from nemo.collections.asr.parts.numba.spec_augment import SpecAugmentNumba, spec_augment_launch_heuristics


Note to myself: remove these as we have a fully vectorized pytorch native implementation that's faster for a while now.

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

# Conflicts: # README.md

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Signed-off-by: artbataev <artbataev@users.noreply.github.com>

artbataev · 2026-03-22T05:38:16Z

/claude review

nemo/core/utils/optional_libs.py

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

artbataev · 2026-03-22T05:43:18Z

/claude review

nemo/core/utils/optional_libs.py

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

artbataev · 2026-03-22T05:48:10Z

/claude review

claude

LGTM

github-actions · 2026-03-22T08:00:39Z

[🤖]: Hi @artbataev 👋,

We wanted to let you know that a CICD pipeline for this PR just finished successfully.

So it might be time to merge this PR or get some approvals.

pzelasko

Excellent work @artbataev!

* Fix int type in grid * Adjust numba and cuda-python requirements * Add cuda-python guards to optional libs * Use cuda-python guards * Add guards for numba Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> --------- Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

…15540) * Fix int type in grid * Adjust numba and cuda-python requirements * Add cuda-python guards to optional libs * Use cuda-python guards * Add guards for numba --------- Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

artbataev added 2 commits March 17, 2026 03:34

Relax numba-cuda version

5cc2d78

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Fix int type in grid

b87a59c

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

artbataev requested review from nithinraok and pzelasko March 17, 2026 11:11

github-actions bot added the ASR label Mar 17, 2026

artbataev added the Run CICD label Mar 17, 2026

artbataev temporarily deployed to test March 17, 2026 11:18 — with GitHub Actions Inactive

github-actions bot removed the Run CICD label Mar 17, 2026

pzelasko reviewed Mar 17, 2026

View reviewed changes

nithinraok reviewed Mar 17, 2026

View reviewed changes

artbataev added 2 commits March 20, 2026 13:27

Merge branch 'main' into vbataev/fix_numba

a6449f3

Adjust numba and cuda-python requirements

cc8ca6f

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

github-actions bot added the core Changes to NeMo Core label Mar 20, 2026

artbataev added 2 commits March 20, 2026 03:05

Add cuda-python guards to optional libs

7af78ec

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Use cuda-python guards

4ad67eb

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

artbataev changed the title ~~Relax numba-cuda version~~ Fix numba-cuda and cuda-python installation and usage Mar 20, 2026

artbataev changed the title ~~Fix numba-cuda and cuda-python installation and usage~~ Fix numba-cuda and cuda-python installation and usage Mar 20, 2026

artbataev and others added 4 commits March 20, 2026 10:06

Apply isort and black reformatting

1e4aa7f

Signed-off-by: artbataev <artbataev@users.noreply.github.com>

Fix flake8 suggestions

e033adf

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Fix flake8 suggestions

05a84a5

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Bugfix

f35ea8c

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

pzelasko reviewed Mar 20, 2026

View reviewed changes

artbataev added 2 commits March 20, 2026 08:48

Add guards for numba

ef5ac48

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Clean up requirements

23d842d

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

pzelasko reviewed Mar 20, 2026

View reviewed changes

Adjust README.md

b64a39f

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

artbataev added the Run CICD label Mar 20, 2026

artbataev added 2 commits March 20, 2026 09:09

Merge remote-tracking branch 'origin/main' into vbataev/fix_numba

5ec7d7b

# Conflicts: # README.md

Adjust README.md

6a1ef80

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

chtruong814 removed the Run CICD label Mar 20, 2026

artbataev added the Run CICD label Mar 22, 2026

Apply isort and black reformatting

396b7d4

Signed-off-by: artbataev <artbataev@users.noreply.github.com>

chtruong814 added Run CICD and removed Run CICD labels Mar 22, 2026

chtruong814 had a problem deploying to test March 22, 2026 05:39 — with GitHub Actions Error

claude bot reviewed Mar 22, 2026

View reviewed changes

nemo/core/utils/optional_libs.py Outdated Show resolved Hide resolved

Fix required version

f055ae6

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

chtruong814 added Run CICD and removed Run CICD labels Mar 22, 2026

chtruong814 had a problem deploying to test March 22, 2026 05:44 — with GitHub Actions Error

claude bot reviewed Mar 22, 2026

View reviewed changes

nemo/core/utils/optional_libs.py Outdated Show resolved Hide resolved

Fix required version

578d37c

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

chtruong814 added Run CICD and removed Run CICD labels Mar 22, 2026

chtruong814 temporarily deployed to test March 22, 2026 05:49 — with GitHub Actions Inactive

claude bot reviewed Mar 22, 2026

View reviewed changes

artbataev enabled auto-merge (squash) March 22, 2026 05:58

github-actions bot removed the Run CICD label Mar 22, 2026

pzelasko approved these changes Mar 23, 2026

View reviewed changes

artbataev merged commit 603e922 into main Mar 23, 2026
132 checks passed

artbataev deleted the vbataev/fix_numba branch March 23, 2026 12:55

This was referenced Mar 23, 2026

[NeMo 2.6] RNNT ASR inference fails on A100 (CUDA 12.8, PyTorch 2.9) with CUDA Graphs error CUDA failure! 35 #15145

Open

Fix building docs #15545

Open

chtruong814 mentioned this pull request Mar 25, 2026

docs: Fix docs build by setting uv conflicts for cu12 vs cu13 #15548

Merged

8 tasks

Conversation

artbataev commented Mar 17, 2026

What does this PR do ?

Changelog

Usage

GitHub Actions CI

Before your PR is "Ready for review"

Who can review?

Additional Information

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nithinraok left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

artbataev Mar 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

artbataev commented Mar 22, 2026

Uh oh!

Uh oh!

artbataev commented Mar 22, 2026

Uh oh!

Uh oh!

artbataev commented Mar 22, 2026

Uh oh!

claude bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Mar 22, 2026

Uh oh!

pzelasko left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

artbataev Mar 20, 2026 •

edited

Loading