chore: import upstream snapshot with attribution

nemo/core/utils/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.core.utils import numba_utils, process_launcher

nemo/core/utils/cuda_python_utils.py

@@ -0,0 +1,275 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import contextlib
+
+import numpy as np
+import torch
+from packaging.version import Version
+
+from nemo.core.utils.optional_libs import CUDA_PYTHON_AVAILABLE, cuda_python_required
+from nemo.utils.exceptions import NeMoBaseException
+
+if CUDA_PYTHON_AVAILABLE:
+    from cuda.bindings import __version__ as cuda_python_version
+    from cuda.bindings import driver as cuda
+    from cuda.bindings import nvrtc
+    from cuda.bindings import runtime as cudart
+
+__CUDA_PYTHON_MINIMUM_VERSION_CUDA_GRAPH_CONDITIONAL_NODES_SUPPORTED__ = (12, 6)  # 12060
+
+
+class NeMoCUDAPythonException(NeMoBaseException):
+    """Exception caused by python-cuda in NeMo"""
+
+    pass
+
+
+CUDA_GRAPH_COMPILE_ERROR_TYPES = (NeMoCUDAPythonException,)
+if hasattr(torch, "AcceleratorError"):
+    CUDA_GRAPH_COMPILE_ERROR_TYPES += (torch.AcceleratorError,)
+
+
+def check_cuda_python_cuda_graphs_conditional_nodes_supported():
+    """Check if CUDA and CUDA-Python are available with CUDA Graphs with conditional nodes support"""
+    # for CPU-only environment we need to raise an exception, otherwise cuda-python library will fail
+    if not torch.cuda.is_available():
+        raise EnvironmentError("CUDA is not available")
+
+    try:
+        from cuda.bindings import driver as cuda
+    except ImportError:
+        raise ModuleNotFoundError("No `cuda-python` module. Please do `pip install cuda-python>=12.3`")
+
+    from cuda.bindings import __version__ as cuda_python_version
+
+    if Version(cuda_python_version) < Version("12.3.0"):
+        raise ImportError(f"Found cuda-python {cuda_python_version}, but at least version 12.3.0 is needed.")
+
+    error, driver_version = cuda.cuDriverGetVersion()
+    if error != cuda.CUresult.CUDA_SUCCESS:
+        raise ImportError(f"cuDriverGetVersion() returned {cuda.cuGetErrorString(error)}")
+
+    driver_version_major = driver_version // 1000
+    driver_version_minor = (driver_version % 1000) // 10
+
+    driver_version = (driver_version_major, driver_version_minor)
+    if driver_version < __CUDA_PYTHON_MINIMUM_VERSION_CUDA_GRAPH_CONDITIONAL_NODES_SUPPORTED__:
+        required_version = __CUDA_PYTHON_MINIMUM_VERSION_CUDA_GRAPH_CONDITIONAL_NODES_SUPPORTED__
+        raise ImportError(
+            f"""Driver supports cuda toolkit version \
+{driver_version_major}.{driver_version_minor}, but the driver needs to support \
+at least {required_version[0]},{required_version[1]}. Please update your cuda driver."""
+        )
+
+
+def skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported():
+    """
+    Helper method to skip pytest test case if cuda graph conditionals nodes are not supported.
+    """
+    try:
+        check_cuda_python_cuda_graphs_conditional_nodes_supported()
+    except (ImportError, ModuleNotFoundError, EnvironmentError) as e:
+        import pytest
+
+        pytest.skip(
+            "Test using cuda graphs with conditional nodes is being skipped because "
+            f"cuda graphs with conditional nodes aren't supported. Error message: {e}"
+        )
+
+
+@cuda_python_required
+def assert_drv(err):
+    """
+    Throws an exception if the return value of a cuda-python call is not success.
+    """
+    if isinstance(err, cuda.CUresult):
+        if err != cuda.CUresult.CUDA_SUCCESS:
+            raise NeMoCUDAPythonException("Cuda Error: {}".format(err))
+    elif isinstance(err, nvrtc.nvrtcResult):
+        if err != nvrtc.nvrtcResult.NVRTC_SUCCESS:
+            raise NeMoCUDAPythonException("Nvrtc Error: {}".format(err))
+    elif isinstance(err, cudart.cudaError_t):
+        if err != cudart.cudaError_t.cudaSuccess:
+            raise NeMoCUDAPythonException("Cuda Runtime Error: {}".format(err))
+    else:
+        raise NeMoCUDAPythonException("Unknown error type: {}".format(err))
+
+
+@cuda_python_required
+def cu_call(f_call_out):
+    """
+    Makes calls to cuda-python's functions inside cuda.cuda more python by throwing an exception
+    if they return a status which is not cudaSuccess
+    """
+    error, *others = f_call_out
+    if error != cudart.cudaError_t.cudaSuccess:
+        raise NeMoCUDAPythonException(f"CUDA failure! {error}")
+    else:
+        return tuple(others)
+
+
+@contextlib.contextmanager
+@cuda_python_required
+def with_conditional_node(while_loop_kernel, while_loop_args, while_loop_conditional_handle, device):
+    """
+    Even though we add a conditional node only once, we need to
+    capture the kernel that calls cudaGraphSetConditional() both
+    before in the parent graph containing the while loop body graph
+    and after the rest of the while loop body graph (because we need
+    to decide both whether to enter the loop, and also whether to
+    execute the next iteration of the loop).
+    """
+    # NB: depending on cuda-python version, cudaStreamGetCaptureInfo can return either 5 or 6 elements
+    capture_status, _, graph, *_ = cu_call(
+        cudart.cudaStreamGetCaptureInfo(torch.cuda.current_stream(device=device).cuda_stream)
+    )
+    assert capture_status == cudart.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive
+
+    cuda.cuLaunchKernel(
+        while_loop_kernel,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        0,
+        torch.cuda.current_stream(device=device).cuda_stream,
+        while_loop_args.ctypes.data,
+        0,
+    )
+
+    # NB: depending on cuda-python version, cudaStreamGetCaptureInfo can return either 5 or 6 elements
+    capture_status, _, graph, dependencies, *_ = cu_call(
+        cudart.cudaStreamGetCaptureInfo(torch.cuda.current_stream(device=device).cuda_stream)
+    )
+    assert capture_status == cudart.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive
+
+    driver_params = cuda.CUgraphNodeParams()
+    driver_params.type = cuda.CUgraphNodeType.CU_GRAPH_NODE_TYPE_CONDITIONAL
+    driver_params.conditional.handle = while_loop_conditional_handle
+    driver_params.conditional.type = cuda.CUgraphConditionalNodeType.CU_GRAPH_COND_TYPE_WHILE
+    driver_params.conditional.size = 1
+    if Version(cuda_python_version) == Version("12.3.0"):
+        # Work around for https://github.com/NVIDIA/cuda-python/issues/55
+        # Originally, cuda-python version 12.3.0 failed to allocate phGraph_out
+        # on its own.
+        # This bug is fixed in cuda-python version 12.4.0. In fact, we can
+        # no longer write to phGraph_out in cuda-python 12.4.0, so we must
+        # condition on the version number.
+        driver_params.conditional.phGraph_out = [cuda.CUgraph()]
+    (ctx,) = cu_call(cuda.cuCtxGetCurrent())
+    driver_params.conditional.ctx = ctx
+
+    # Use driver API here because of bug in cuda-python runtime API: https://github.com/NVIDIA/cuda-python/issues/55
+    # TODO: Change call to this after fix goes in (and we bump minimum cuda-python version to 12.4.0):
+    # node, = cu_call(cudart.cudaGraphAddNode(graph, dependencies, len(dependencies), driver_params))
+    # CUDA 13 (cuda-python >= 13.0.0) adds an edgeData parameter to cuGraphAddNode and
+    # cudaStreamUpdateCaptureDependencies; CUDA 12 does not accept it.
+    _cuda13 = Version(cuda_python_version) >= Version("13.0.0")
+    if _cuda13:
+        (node,) = cu_call(cuda.cuGraphAddNode(graph, dependencies, None, len(dependencies), driver_params))
+    else:
+        (node,) = cu_call(cuda.cuGraphAddNode(graph, dependencies, len(dependencies), driver_params))
+    body_graph = driver_params.conditional.phGraph_out[0]
+
+    if _cuda13:
+        cu_call(
+            cudart.cudaStreamUpdateCaptureDependencies(
+                torch.cuda.current_stream(device=device).cuda_stream,
+                [node],
+                None,
+                1,
+                cudart.cudaStreamUpdateCaptureDependenciesFlags.cudaStreamSetCaptureDependencies,
+            )
+        )
+    else:
+        cu_call(
+            cudart.cudaStreamUpdateCaptureDependencies(
+                torch.cuda.current_stream(device=device).cuda_stream,
+                [node],
+                1,
+                cudart.cudaStreamUpdateCaptureDependenciesFlags.cudaStreamSetCaptureDependencies,
+            )
+        )
+    body_stream = torch.cuda.Stream(device)
+    previous_stream = torch.cuda.current_stream(device=device)
+    body_capture_active = False
+    cu_call(
+        cudart.cudaStreamBeginCaptureToGraph(
+            body_stream.cuda_stream,
+            body_graph,
+            None,
+            None,
+            0,
+            cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal,
+        )
+    )
+    body_capture_active = True
+
+    try:
+        torch.cuda.set_stream(body_stream)
+        yield body_stream, body_graph
+
+        cuda.cuLaunchKernel(
+            while_loop_kernel, 1, 1, 1, 1, 1, 1, 0, body_stream.cuda_stream, while_loop_args.ctypes.data, 0
+        )
+
+        end_capture_out = cudart.cudaStreamEndCapture(body_stream.cuda_stream)
+        body_capture_active = False
+        cu_call(end_capture_out)
+    finally:
+        if body_capture_active:
+            try:
+                end_capture_out = cudart.cudaStreamEndCapture(body_stream.cuda_stream)
+                body_capture_active = False
+                cu_call(end_capture_out)
+            except Exception:
+                pass
+        torch.cuda.set_stream(previous_stream)
+
+
+@cuda_python_required
+def run_nvrtc(kernel_string: str, kernel_name: bytes, program_name: bytes):
+    """Run CUDA kernel using CUDA-Python"""
+    err, prog = nvrtc.nvrtcCreateProgram(str.encode(kernel_string), program_name, 0, [], [])
+    assert_drv(err)
+    # Compile program
+    # Not specifying --gpu-architecture will default us to a fairly low compute capability, which is a safe bet.
+    # Otherwise, there are ways to query the current device's compute capability.
+    # https://stackoverflow.com/questions/48283009/nvcc-get-device-compute-capability-in-runtime
+    opts = []
+    (err,) = nvrtc.nvrtcCompileProgram(prog, len(opts), opts)
+    assert_drv(err)
+    err, size = nvrtc.nvrtcGetProgramLogSize(prog)
+    assert_drv(err)
+    buf = b" " * size
+    (err,) = nvrtc.nvrtcGetProgramLog(prog, buf)
+    assert_drv(err)
+
+    # Get PTX from compilation
+    err, ptxSize = nvrtc.nvrtcGetPTXSize(prog)
+    assert_drv(err)
+    ptx = b" " * ptxSize
+    (err,) = nvrtc.nvrtcGetPTX(prog, ptx)
+    assert_drv(err)
+
+    ptx = np.char.array(ptx)
+    err, module = cuda.cuModuleLoadData(ptx.ctypes.data)
+    assert_drv(err)
+    err, kernel = cuda.cuModuleGetFunction(module, kernel_name)
+    assert_drv(err)
+
+    return kernel

nemo/core/utils/k2_guard.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Guard for importing optional NeMo dependency k2.
+Contains checks for k2 availability and version.
+Use `from nemo.core.utils.k2_guard import k2` to import k2 instead of direct import.
+If there is an error, the module will raise an exception with a helpful message.
+"""
+
+import textwrap
+
+from lightning.pytorch.utilities.imports import package_available
+from packaging.version import Version
+
+from nemo.core.utils.k2_utils import K2_INSTALLATION_MESSAGE
+
+__K2_MINIMUM_MAJOR_VERSION = 1
+__K2_MINIMUM_MINOR_VERSION = 14
+
+__K2_MINIMUM_VERSION = Version(f"{__K2_MINIMUM_MAJOR_VERSION}.{__K2_MINIMUM_MINOR_VERSION}")
+
+
+if not package_available("k2"):
+    raise ModuleNotFoundError("Module k2 is not available.\n" + K2_INSTALLATION_MESSAGE)
+
+import k2  # noqa: E402
+
+try:
+    __k2_version = Version(k2.__dev_version__)
+except AttributeError:
+    raise ImportError("Module k2 is corrupted.\n" + K2_INSTALLATION_MESSAGE)
+
+if __k2_version < __K2_MINIMUM_VERSION:
+    raise ImportError(
+        textwrap.dedent(
+            f"""
+            Minimum required k2 version: {__K2_MINIMUM_VERSION};
+            Installed k2 version: {__k2_version}
+            """
+        )
+    )

nemo/core/utils/k2_utils.py

@@ -0,0 +1,24 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+K2_INSTALLATION_MESSAGE = (
+    "Could not import `k2`.\n"
+    "Please install k2 in one of the following ways:\n"
+    "1) (recommended) Run `bash scripts/installers/install_k2.sh`\n"
+    "2) Use any approach from https://k2-fsa.github.io/k2/installation/index.html "
+    "if your your cuda and pytorch versions are supported.\n"
+    "It is advised to always install k2 using setup.sh only, "
+    "as different versions of k2 may not interact with the NeMo code as expected."
+)

nemo/core/utils/neural_type_utils.py

@@ -0,0 +1,82 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections import defaultdict
+from typing import Dict, List, Optional
+import torch
+from nemo.core.neural_types import AxisKind, NeuralType
+
+
+def get_io_names(types: Optional[Dict[str, NeuralType]], disabled_names: List[str]) -> List[str]:
+    if types is None:
+        return []
+    names = list(types.keys())
+    for name in disabled_names:
+        if name in names:
+            names.remove(name)
+    return names
+
+
+def extract_dynamic_axes(name: str, ntype: NeuralType):
+    """
+    This method will extract BATCH and TIME dimension ids from each provided input/output name argument.
+
+    For example, if module/model accepts argument named "input_signal" with type corresponding to [Batch, Time, Dim]
+    shape, then the returned result should contain "input_signal" -> [0, 1] because Batch and Time are dynamic axes
+    as they can change from call to call during inference.
+
+    Args:
+        name: Name of input or output parameter
+        ntype: Corresponding Neural Type
+
+    Returns:
+
+    """
+
+    def unpack_nested_neural_type(neural_type):
+        if type(neural_type) in (list, tuple):
+            return unpack_nested_neural_type(neural_type[0])
+        return neural_type
+
+    dynamic_axes = defaultdict(list)
+    if type(ntype) in (list, tuple):
+        ntype = unpack_nested_neural_type(ntype)
+
+    if ntype.axes:
+        for ind, axis in enumerate(ntype.axes):
+            if axis.kind in [AxisKind.Batch, AxisKind.Time, AxisKind.Width, AxisKind.Height]:
+                dynamic_axes[name].append(ind)
+    return dynamic_axes
+
+
+def get_dynamic_axes(types, names, use_dynamo=False):
+    dynamic_axes = defaultdict(list)
+    if names is not None:
+        for name in names:
+            if name in types:
+                dynamic_axes.update(extract_dynamic_axes(name, types[name]))
+    if use_dynamo:
+        dynamic_shapes = {}
+        batch = torch.export.Dim("batch")
+        for name, dims in dynamic_axes.items():
+            ds = {}
+            for d in dims:
+                if d == 0:
+                    ds[d] = batch
+                # this currently has issues: https://github.com/pytorch/pytorch/issues/126127
+                else:
+                    ds[d] = torch.export.Dim(name + '__' + str(d))
+            dynamic_shapes[name] = ds
+        dynamic_axes = dynamic_shapes
+    return dynamic_axes

nemo/core/utils/numba_utils.py

@@ -0,0 +1,195 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import contextlib
+import logging as pylogger
+import operator
+import os
+
+from typing import Tuple, Union
+
+from nemo.utils import model_utils
+
+# Prevent Numba CUDA logs from showing at info level
+cuda_logger = pylogger.getLogger('numba.cuda.cudadrv.driver')
+cuda_logger.setLevel(pylogger.ERROR)  # only show error
+
+__NUMBA_DEFAULT_MINIMUM_VERSION__ = "0.53.0"
+__NUMBA_MINIMUM_VERSION__ = os.environ.get("NEMO_NUMBA_MINVER", __NUMBA_DEFAULT_MINIMUM_VERSION__)
+
+__NUMBA_MINIMUM_VERSION_FP16_SUPPORTED__ = "0.57.0"
+
+
+NUMBA_INSTALLATION_MESSAGE = (
+    "Could not import `numba`.\n"
+    "Please install numba in one of the following ways."
+    "1) If using conda, simply install it with conda using `conda install -c numba numba`\n"
+    "2) If using pip (not recommended), `pip install --upgrade numba`\n"
+    "followed by `export NUMBAPRO_LIBDEVICE='/usr/local/cuda/nvvm/libdevice/'` and \n"
+    "`export NUMBAPRO_NVVM='/usr/local/cuda/nvvm/lib64/libnvvm.so'`.\n"
+    "It is advised to always install numba using conda only, "
+    "as pip installations might interfere with other libraries such as llvmlite.\n"
+    "If pip install does not work, you can also try adding `--ignore-installed` to the pip command,\n"
+    "but this is not advised."
+)
+
+STRICT_NUMBA_COMPAT_CHECK = True
+
+# Get environment key if available
+if 'STRICT_NUMBA_COMPAT_CHECK' in os.environ:
+    check_str = os.environ.get('STRICT_NUMBA_COMPAT_CHECK')
+    check_bool = str(check_str).lower() in ("yes", "true", "t", "1")
+    STRICT_NUMBA_COMPAT_CHECK = check_bool
+
+
+def is_numba_compat_strict() -> bool:
+    """
+    Returns strictness level of numba cuda compatibility checks.
+
+    If value is true, numba cuda compatibility matrix must be satisfied.
+    If value is false, only cuda availability is checked, not compatibility.
+    Numba Cuda may still compile and run without issues in such a case, or it may fail.
+    """
+    return STRICT_NUMBA_COMPAT_CHECK
+
+
+def set_numba_compat_strictness(strict: bool):
+    """
+    Sets the strictness level of numba cuda compatibility checks.
+
+    If value is true, numba cuda compatibility matrix must be satisfied.
+    If value is false, only cuda availability is checked, not compatibility.
+    Numba Cuda may still compile and run without issues in such a case, or it may fail.
+
+    Args:
+        strict: bool value, whether to enforce strict compatibility checks or relax them.
+    """
+    global STRICT_NUMBA_COMPAT_CHECK
+    STRICT_NUMBA_COMPAT_CHECK = strict
+
+
+@contextlib.contextmanager
+def with_numba_compat_strictness(strict: bool):
+    """Context manager for setting numba compatibility checks temporary"""
+    initial_strictness = is_numba_compat_strict()
+    set_numba_compat_strictness(strict=strict)
+    yield
+    set_numba_compat_strictness(strict=initial_strictness)
+
+
+def numba_cpu_is_supported(min_version: str) -> bool:
+    """
+    Tests if an appropriate version of numba is installed.
+
+    Args:
+        min_version: The minimum version of numba that is required.
+
+    Returns:
+        bool, whether numba CPU supported with this current installation or not.
+    """
+    module_available, msg = model_utils.check_lib_version('numba', checked_version=min_version, operator=operator.ge)
+
+    # If numba is not installed
+    if module_available is None:
+        return False
+    else:
+        return True
+
+
+def numba_cuda_is_supported(min_version: str) -> bool:
+    """
+    Tests if an appropriate version of numba is installed, and if it is,
+    if cuda is supported properly within it.
+
+    Args:
+        min_version: The minimum version of numba that is required.
+
+    Returns:
+        bool, whether cuda is supported with this current installation or not.
+    """
+    module_available = numba_cpu_is_supported(min_version)
+
+    # If numba is not installed
+    if module_available is None:
+        return False
+
+    # If numba version is installed and available
+    if module_available is True:
+        from numba import cuda
+
+        try:
+            cuda_available = cuda.is_available()
+            if cuda_available:
+                cuda_compatible = cuda.cudadrv.runtime.get_version()[0] in (12, 13)
+            else:
+                cuda_compatible = False
+
+            if is_numba_compat_strict():
+                return cuda_available and cuda_compatible
+            else:
+                return cuda_available
+
+        except Exception:
+            # dlopen(libcudart.dylib) might fail if CUDA was never installed in the first place.
+            return False
+
+    else:
+        return False
+
+
+def is_numba_cuda_fp16_supported(return_reason: bool = False) -> Union[bool, Tuple[bool, str]]:
+    """
+    Utility method that returns a bool, stating if FP16 is supported for numba cuda kernels or not.
+
+    Returns:
+        bool, whether Numba CUDA will support fp16 or not.
+    """
+    reason = ""
+    use_nvidia_binding = os.environ.get('NUMBA_CUDA_USE_NVIDIA_BINDING', None)
+    if use_nvidia_binding is not None:
+        use_nvidia_binding = use_nvidia_binding.lower() == "1"
+        reason += "Env variable `NUMBA_CUDA_USE_NVIDIA_BINDING` is available and set to `1`. "
+    else:
+        use_nvidia_binding = False
+        reason += "Env variable `NUMBA_CUDA_USE_NVIDIA_BINDING` is not available or has not set to `1`."
+
+    numba_fp16_version_correct = model_utils.check_lib_version(
+        'numba', __NUMBA_MINIMUM_VERSION_FP16_SUPPORTED__, operator=operator.ge
+    )[0]
+
+    if numba_fp16_version_correct:
+        reason += "Numba CUDA FP16 is supported in installed numba version."
+    else:
+        reason += "Numba CUDA FP16 is not supported in installed numba version."
+
+    result = use_nvidia_binding and numba_fp16_version_correct
+
+    if return_reason:
+        return result, reason
+    else:
+        return result
+
+
+def skip_numba_cuda_test_if_unsupported(min_version: str):
+    """
+    Helper method to skip pytest test case if numba cuda is not supported.
+
+    Args:
+        min_version: The minimum version of numba that is required.
+    """
+    numba_cuda_support = numba_cuda_is_supported(min_version)
+    if not numba_cuda_support:
+        import pytest
+
+        pytest.skip(f"Numba cuda test is being skipped. Minimum version required : {min_version}")

nemo/core/utils/optional_libs.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Module with guards for optional libraries, that cannot be listed in `requirements.txt`.
+Provides helper constants and decorators to check if the library is available in the system.
+"""
+
+__all__ = [
+    # kenlm
+    "KENLM_AVAILABLE",
+    "kenlm_required",
+    # k2
+    "K2_AVAILABLE",
+    "k2_required",
+    # triton
+    "TRITON_AVAILABLE",
+    "triton_required",
+    # cuda-python
+    "CUDA_PYTHON_AVAILABLE",
+    "cuda_python_required",
+    # numba
+    "NUMBA_AVAILABLE",
+    "numba_required",
+    # numba-cuda
+    "NUMBA_CUDA_AVAILABLE",
+    "numba_cuda_required",
+]
+
+import importlib.util
+from functools import wraps
+
+from packaging.version import Version
+
+from nemo.core.utils.k2_utils import K2_INSTALLATION_MESSAGE
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__, numba_cpu_is_supported, numba_cuda_is_supported
+
+
+def is_lib_available(name: str) -> bool:
+    """
+    Checks if the library/package with `name` is available in the system
+    NB: try/catch with importlib.import_module(name) requires importing the library, which can be slow.
+    So, `find_spec` should be preferred
+    """
+    return importlib.util.find_spec(name) is not None
+
+
+KENLM_AVAILABLE = is_lib_available("kenlm")
+KENLM_INSTALLATION_MESSAGE = "Try installing kenlm with `pip install kenlm`"
+
+TRITON_AVAILABLE = is_lib_available("triton")
+TRITON_INSTALLATION_MESSAGE = "Try installing triton with `pip install triton`"
+
+NUMBA_AVAILABLE = numba_cpu_is_supported(__NUMBA_MINIMUM_VERSION__)
+NUMBA_INSTALLATION_MESSAGE = (
+    "Numba is not found. Install with `pip install numba`. "
+    "For GPU support install with `pip install numba-cuda[cu12]` or `pip install numba-cuda[cu13]`"
+)
+
+NUMBA_CUDA_AVAILABLE = numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)
+NUMBA_CUDA_INSTALLATION_MESSAGE = (
+    "Numba with GPU support is not available. "
+    "For GPU support install with `pip install numba-cuda[cu12]` or `pip install numba-cuda[cu13]`"
+)
+
+try:
+    from nemo.core.utils.k2_guard import k2 as _  # noqa: F401
+
+    K2_AVAILABLE = True
+except (ImportError, ModuleNotFoundError):
+    K2_AVAILABLE = False
+
+try:
+    from cuda.bindings import __version__ as cuda_python_version
+
+    if Version(cuda_python_version) >= Version("12.6.0"):
+        CUDA_PYTHON_AVAILABLE = True
+    else:
+        CUDA_PYTHON_AVAILABLE = False
+except (ImportError, ModuleNotFoundError):
+    CUDA_PYTHON_AVAILABLE = False
+
+CUDA_PYTHON_INSTALLATION_MESSAGE = "Try installing cuda-python with `pip install cuda-python>=12.6.0`"
+
+
+def identity_decorator(f):
+    """Identity decorator for further using in conditional decorators"""
+    return f
+
+
+def _lib_required(is_available: bool, name: str, message: str | None = None):
+    """
+    Decorator factory. Returns identity decorator if lib `is_available`,
+    otherwise returns a decorator which returns a function that raises an error when called.
+    Such decorator can be used for conditional checks for optional libraries in functions and methods
+    with zero computational overhead.
+    """
+    if is_available:
+        return identity_decorator
+
+    # return wrapper that will raise an error when the function is called
+    def function_stub_with_error_decorator(f):
+        """Decorator that replaces the function and raises an error when called"""
+
+        @wraps(f)
+        def wrapper(*args, **kwargs):
+            error_msg = f"Module {name} required for the function {f.__name__} is not found."
+            if message:
+                error_msg += f" {message}"
+            raise ModuleNotFoundError(error_msg)
+
+        return wrapper
+
+    return function_stub_with_error_decorator
+
+
+kenlm_required = _lib_required(is_available=KENLM_AVAILABLE, name="kenlm", message=KENLM_INSTALLATION_MESSAGE)
+triton_required = _lib_required(is_available=TRITON_AVAILABLE, name="triton", message=TRITON_INSTALLATION_MESSAGE)
+k2_required = _lib_required(is_available=K2_AVAILABLE, name="k2", message=K2_INSTALLATION_MESSAGE)
+cuda_python_required = _lib_required(
+    is_available=CUDA_PYTHON_AVAILABLE, name="cuda_python", message=CUDA_PYTHON_INSTALLATION_MESSAGE
+)
+numba_required = _lib_required(is_available=NUMBA_AVAILABLE, name="numba", message=NUMBA_INSTALLATION_MESSAGE)
+numba_cuda_required = _lib_required(
+    is_available=NUMBA_CUDA_AVAILABLE, name="numba-cuda", message=NUMBA_CUDA_INSTALLATION_MESSAGE
+)

nemo/core/utils/process_launcher/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.core.utils.process_launcher.launcher import ProcessLauncher, ProcessLauncherConfig

nemo/core/utils/process_launcher/launcher.py

@@ -0,0 +1,386 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import hashlib
+import os
+import subprocess
+import sys
+import threading
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Optional, Sequence
+
+import torch
+from hydra.core.config_store import ConfigStore
+from hydra.core.hydra_config import HydraConfig
+from hydra.core.plugins import Plugins
+from hydra.core.singleton import Singleton
+from hydra.core.utils import JobReturn, JobStatus, configure_log, filter_overrides, setup_globals
+from hydra.plugins.launcher import Launcher
+from hydra.types import HydraContext, TaskFunction
+from omegaconf import DictConfig, OmegaConf, open_dict
+
+from nemo.utils import logging
+
+
+# monkey-patch hydra func
+def is_in_toplevel_plugins_module(*args, **kwargs) -> bool:
+    """Treat NeMo launcher plugins as top-level Hydra plugins."""
+    return True
+
+
+# Monkey-patch Hydra
+Plugins.instance().is_in_toplevel_plugins_module = is_in_toplevel_plugins_module
+
+
+@dataclass
+class ProcessLauncherConfig:
+    """Configuration for the NeMo process launcher."""
+
+    _target_: str = "nemo.core.utils.process_launcher.launcher.ProcessLauncher"
+    num_gpus: int = -1
+    jobs_per_gpu: int = 1
+
+
+def execute_job(
+    idx: int,
+    overrides: Sequence[str],
+    hydra_context: HydraContext,
+    config: DictConfig,
+    singleton_state: Dict[Any, Any],
+    gpu_idx: int,
+):
+    """
+    Creates a process that launches a "single" job that is identical in config + updated with sweep hyperparams.
+    Since a different process is being used, CUDA can work in non-ddp mode without issue.
+    Attempting ddp when using this script will not work as ddp cannot be used in shared contexts.
+
+    Args:
+        idx: Global index of the job.
+        overrides: List of str overrides that correspond to this job
+        hydra_context: Hydra Context used to load the sweep params into the global config
+        config: Global config that will be updated with sweep hyper parameters.
+        singleton_state: Hydra state.
+        gpu_idx: The GPU ID on which this process will be run.
+
+    Returns:
+        - The Process object that corresponds to this sweep
+        - The JobReturn object holding some metadata about this run
+    """
+    # Required by Hydra (lookup other Hydra Launchers for details)
+    setup_globals()
+    Singleton.set_state(singleton_state)
+
+    # Update base config with overrides to create sweep config
+    sweep_config = hydra_context.config_loader.load_sweep_config(config, list(overrides))
+    with open_dict(sweep_config):
+        sweep_config.hydra.job.id = "{}_{}".format(sweep_config.hydra.job.name, idx)
+        sweep_config.hydra.job.num = idx
+    HydraConfig.instance().set_config(sweep_config)
+
+    # Setup a directory where the config will temporarily be stored.
+    script_path = os.path.join(os.getcwd(), sys.argv[0])
+    script_path = os.path.abspath(script_path)
+
+    hash_salt = "|".join([script_path, str(OmegaConf.to_yaml(config))]).encode('utf-8')
+    hash_val = hashlib.sha256(hash_salt).hexdigest()
+
+    config_dir = os.path.join(os.getcwd(), "hydra_cfg", str(hash_val))
+    if not os.path.exists(config_dir):
+        os.makedirs(config_dir, exist_ok=True)
+
+    task_cfg = copy.deepcopy(sweep_config)
+
+    # Remove hydra from sweep config
+    # This is done to prevent recursive call to multirun launcher in Hydra.
+    with open_dict(task_cfg):
+        task_cfg.pop('hydra', '')
+
+    # Save the current jobs config to directory
+    temp_config_name = f"config_{idx}.yaml"
+    temp_config = os.path.join(config_dir, temp_config_name)
+    OmegaConf.save(task_cfg, temp_config)
+
+    # Compute the overides as a dict
+    overrides = OmegaConf.to_container(config.hydra.overrides.task)
+
+    # Check and replace trainer.devices in config with gpu_idx
+    found_devices = False
+    gpu_override = f'trainer.devices=[{gpu_idx}]'
+    for oidx, val in enumerate(overrides):
+        if 'trainer.devices' in val:
+            overrides[oidx] = gpu_override
+            found_devices = True
+
+    if not found_devices:
+        overrides.append(gpu_override)
+
+    # Build launch command
+    # Note: We depend on PTL doing the right thing since this command has global visibility of all CUDA_VISIBLE_DEVICES
+    cmd = [
+        'python',
+        script_path,
+        "--config-path",
+        config_dir,
+        "--config-name",
+        temp_config_name,
+        *overrides,
+    ]
+
+    # Launch the subprocess; pipe the stderr
+    # NOTE: If this hangs due to some reason after prolonged training, it means that the stderr pipe buffer
+    # has become full at the OS level and we need to explicitly empty it (either parallel thread or manually
+    # call proc.communicate(). It should not happen in general case as stderr is filled only in case retcode != 0
+    # If it does happen though, implement the code here
+    # https://stackoverflow.com/questions/39607172/python-subprocess-popen-poll-seems-to-hang-but-communicate-works
+    proc = subprocess.Popen(cmd, stderr=subprocess.PIPE)
+
+    # Setup data thread for stderr
+    std_error_buffer = []
+    # Trivial thread just reads lines from stdout into the list
+    drainerthread = threading.Thread(target=std_error_buffer.extend, args=(proc.stderr,))
+    drainerthread.daemon = True
+    drainerthread.start()
+
+    # Construct JobReturn object for Hydra
+    res = JobReturn()
+    res.cfg = task_cfg
+    res.overrides = overrides
+    res.hydra_cfg = config
+    res.working_dir = os.getcwd()
+    res.return_value = None
+
+    return proc, res, (std_error_buffer, drainerthread)
+
+
+def launch(
+    launcher,
+    job_overrides: Sequence[Sequence[str]],
+    initial_job_idx: int,
+) -> Sequence[JobReturn]:
+    """
+    Args:
+        launcher: Reference to the Launched subclass
+        job_overrides: A List of List<String>, where each inner list is the arguments for one job run
+        initial_job_idx: Initial job idx in batch
+
+    Returns:
+        A list of JobReturn objects.
+    """
+    # Needed for Hydra, lookup JoblibLauncher in Hydra
+    setup_globals()
+    assert launcher.config is not None
+    assert launcher.task_function is not None
+    assert launcher.hydra_context is not None
+
+    configure_log(launcher.config.hydra.hydra_logging, launcher.config.hydra.verbose)
+    sweep_dir = Path(str(launcher.config.hydra.sweep.dir))
+    sweep_dir.mkdir(parents=True, exist_ok=True)
+
+    # Extraact the runner's config (its actually a DictConfig, but type is used for autocomplete)
+    runner_cfg = launcher.runner  # type: ProcessLauncherConfig
+
+    logging.info(
+        "ProcessLauncher({}) is launching {} jobs".format(
+            ",".join([f"{k}={v}" for k, v in runner_cfg.items()]),
+            len(job_overrides),
+        )
+    )
+    logging.info("Launching jobs, sweep output dir : {}".format(sweep_dir))
+    for idx, overrides in enumerate(job_overrides):
+        logging.info("\t#{} : {}".format(idx, " ".join(filter_overrides(overrides))))
+
+    # Needed by Hydra
+    singleton_state = Singleton.get_state()
+
+    # Process the runner's config to build up the multiplex config
+    num_gpus = runner_cfg.get('num_gpus', -1)
+    jobs_per_gpu = runner_cfg.get('jobs_per_gpu', 1)
+
+    # Only GPUs are supported for now.
+    if num_gpus <= 0:
+        if torch.cuda.is_available():
+            num_gpus = torch.cuda.device_count()
+        else:
+            raise ValueError(f"{launcher.__class__.__name__} only supports GPU operations.")
+
+    # Setup arguments for multiplex runner
+    overrides = list(job_overrides)
+    num_overrides = len(overrides)
+
+    job_idx = 0
+    batch_size = num_gpus * jobs_per_gpu
+    gpu_idx = 0
+
+    ret = []  # List of returned JobResult
+    subprocess_list = []  # Buffer of subprocess
+    results = []  # Buffer of JobResult
+
+    # STD ERROR cache
+    std_error_buffers = []  # type: List[List[str]]
+    std_error_threads = []  # type: threading.Thread
+
+    # Run over all job combinations
+    while job_idx < num_overrides:
+        # Fill up subprocess buffer while its size is smaller than multiplex batch size
+        while len(subprocess_list) < batch_size:
+            # If we run out of jobs, stop trying to submit more jobs
+            if job_idx >= num_overrides:
+                break
+
+            # Submit a job as a new process
+            process, res, error_tup = execute_job(
+                initial_job_idx + job_idx,
+                overrides[job_idx],
+                launcher.hydra_context,
+                launcher.config,
+                singleton_state,
+                gpu_idx % num_gpus,  # This will evenly distribute GPU load
+            )
+
+            # Store the subprocesses and JobResults
+            subprocess_list.append(process)
+            results.append(res)
+
+            # Manage stderror thread data
+            std_error_buffers.append(error_tup[0])
+            std_error_threads.append(error_tup[1])
+
+            job_idx += 1
+            gpu_idx += 1
+
+        # Poll for samples in batch to finish.
+        if len(subprocess_list) > 0:
+            finished_processes = [0] * len(subprocess_list)
+
+            # Check if all processes are completed or not
+            # This is busy waiting, this is actually quite necessary
+            # Turns out that when you do proc.communicate(), you block all other threads immediately.
+            # IE they may fill up their buffers entirely, and hang while they wait for the first thread
+            # who called communicate() to finish its work or crash.
+            # Effectively it entirely stops multiprocessing jobs or multiplexed runs.
+            # Must poll and busy wait to keep threads alive, along with drain the pipes with thread buffers.
+            while sum(finished_processes) < len(subprocess_list):
+                # Check all processes to make sure they have a retcode (doesnt matter yet if 0 or not)
+                for proc_idx, proc in enumerate(subprocess_list):
+                    # poll() is cheaper op than communicate()
+                    retcode = proc.poll()
+
+                    if retcode is not None:
+                        # Log that the process with some ID has finished
+                        if finished_processes[proc_idx] == 0:
+                            logging.info(f"Processed job : {len(ret) + proc_idx} :: Ret code = {retcode}")
+
+                        finished_processes[proc_idx] = 1
+
+                        # Join this thread and merge its stderror buffer
+                        proc.wait()
+                        std_error_threads[proc_idx].join()
+                        error_data = std_error_buffers[proc_idx]
+                        error_data = [
+                            str(data, encoding='utf-8').encode('utf-8').decode('utf-8').encode('utf-8')
+                            for data in error_data
+                        ]
+
+                        std_error_buffers[proc_idx] = error_data
+
+                time.sleep(1.0)
+
+            # Process all the subprocess results
+            for proc_idx, (proc, res) in enumerate(zip(subprocess_list, results)):
+                # Wait until completion of process
+                output, error = proc.communicate()
+
+                # 0 is for successful run
+                if proc.returncode == 0:
+                    res.status = JobStatus.COMPLETED
+                else:
+                    # > 0 is for error, log the error.
+                    # Note: For the sake of efficiency while we log the error and raise an exception,
+                    # It will only raise the 1st wrong job in all the jobs.
+                    # If multiple jobs fail, it will still try to execute every job first before
+                    # raising the error for the first one.
+                    # This is done so that even if some jobs fail (say OOM or something),
+                    # other jobs can still run.
+                    err_buffer = std_error_buffers[proc_idx]
+                    if isinstance(err_buffer, (list, tuple)):
+                        err_string = ""
+                        for err_line in err_buffer:
+                            err_string = (
+                                err_string + f"{str(err_line, encoding='utf-8').encode('utf-8').decode('utf-8')}"
+                            )
+                    else:
+                        err_string = err_buffer
+
+                    error_msg = (
+                        f"\nHyperparameter Arguments : {proc.args}\n"
+                        f"Process Return code : {proc.returncode}\n"
+                        f"Error Trace :\n"
+                        f"{err_string}"
+                    )
+                    res.return_value = Exception(error_msg)
+                    res.status = JobStatus.FAILED
+
+                logging.info(f"Finished executing job : {len(ret)}. Return Code = {proc.returncode}")
+                ret.append(res)
+
+            # Reset for next batch
+            subprocess_list.clear()
+            results.clear()
+
+    return ret
+
+
+class ProcessLauncher(Launcher):
+    """Hydra launcher that multiplexes jobs across local GPU processes."""
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Process Launcher
+        Based on the JoblibLauncher, but uses processes to scatter jobs in a multiplexed manner across
+        some number of GPUs on a single machine.
+        """
+        self.config: Optional[DictConfig] = None
+        self.task_function: Optional[TaskFunction] = None
+        self.hydra_context: Optional[HydraContext] = None
+
+        self.runner = kwargs  # type: ProcessLauncherConfig
+
+    def setup(
+        self,
+        *,
+        hydra_context: HydraContext,
+        task_function: TaskFunction,
+        config: DictConfig,
+    ) -> None:
+        """Store Hydra launch context and task function."""
+        self.config = config
+        self.task_function = task_function
+        self.hydra_context = hydra_context
+
+    def launch(self, job_overrides: Sequence[Sequence[str]], initial_job_idx: int) -> Sequence[JobReturn]:
+        """Launch jobs with the configured process launcher."""
+
+        return launch(launcher=self, job_overrides=job_overrides, initial_job_idx=initial_job_idx)
+
+
+ConfigStore.instance().store(
+    group="hydra/launcher",
+    name="nemo_launcher",
+    node=ProcessLauncherConfig,
+    provider="nemo_process_launcher",
+)
+
+Plugins.instance().register(ProcessLauncher)