Merge remote-tracking branch 'nemo/main' into magpietts_evaluation_parallelization

rfejgin · rfejgin · commit 7ae45bed9b90 · 2026-02-20T21:30:04.000-08:00
diff --git a/nemo/collections/asr/parts/submodules/ctc_batched_beam_decoding.py b/nemo/collections/asr/parts/submodules/ctc_batched_beam_decoding.py
@@ -586,7 +586,8 @@ def _full_graph_compile(self):
             torch.cuda.graph(self.full_graph, stream=stream_for_graph, capture_error_mode="thread_local"),
         ):
             self._before_process_batch()
-            capture_status, _, graph, _, _, _ = cu_call(
+            # 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=self.state.device).cuda_stream)
             )
 
diff --git a/nemo/collections/asr/parts/submodules/ctc_greedy_decoding.py b/nemo/collections/asr/parts/submodules/ctc_greedy_decoding.py
@@ -874,7 +874,8 @@ def _full_graph_compile(self):
         ):
             self._before_loop()
 
-            capture_status, _, graph, _, _, _ = cu_call(
+            # 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=self.state.device).cuda_stream)
             )
             assert capture_status == cudart.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive
diff --git a/nemo/collections/asr/parts/submodules/cuda_graph_rnnt_greedy_decoding.py b/nemo/collections/asr/parts/submodules/cuda_graph_rnnt_greedy_decoding.py
@@ -205,7 +205,8 @@ def _reinitialize(self, max_time, batch_size, encoder_output, encoder_output_len
             # Get max sequence length
             self.max_out_len_t = self.encoder_output_length.max()
 
-            capture_status, _, graph, _, _, _ = cu_call(
+            # 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=self.device).cuda_stream)
             )
             assert capture_status == cudart.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive
diff --git a/nemo/collections/asr/parts/submodules/rnnt_malsd_batched_computer.py b/nemo/collections/asr/parts/submodules/rnnt_malsd_batched_computer.py
@@ -906,7 +906,8 @@ def _full_graph_compile(self):
             torch.cuda.graph(self.full_graph, stream=stream_for_graph, capture_error_mode="thread_local"),
         ):
             self._before_loop()
-            capture_status, _, graph, _, _, _ = cu_call(
+            # 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=self.state.device).cuda_stream)
             )
 
diff --git a/nemo/collections/asr/parts/submodules/tdt_malsd_batched_computer.py b/nemo/collections/asr/parts/submodules/tdt_malsd_batched_computer.py
@@ -987,7 +987,8 @@ def _full_graph_compile(self):
             torch.cuda.graph(self.full_graph, stream=stream_for_graph, capture_error_mode="thread_local"),
         ):
             self._before_loop()
-            capture_status, _, graph, _, _, _ = cu_call(
+            # 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=self.state.device).cuda_stream)
             )
 
diff --git a/nemo/collections/asr/parts/submodules/transducer_decoding/rnnt_label_looping.py b/nemo/collections/asr/parts/submodules/transducer_decoding/rnnt_label_looping.py
@@ -968,7 +968,8 @@ def _full_graph_compile(self):
         ):
             self._before_outer_loop()
 
-            capture_status, _, graph, _, _, _ = cu_call(
+            # 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=self.state.device).cuda_stream)
             )
             assert capture_status == cudart.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive
diff --git a/nemo/collections/asr/parts/submodules/transducer_decoding/tdt_label_looping.py b/nemo/collections/asr/parts/submodules/transducer_decoding/tdt_label_looping.py
@@ -1041,7 +1041,8 @@ def _full_graph_compile(self):
         ):
             self._before_outer_loop()
 
-            capture_status, _, graph, _, _, _ = cu_call(
+            # 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=self.state.device).cuda_stream)
             )
             assert capture_status == cudart.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive
diff --git a/nemo/core/utils/cuda_python_utils.py b/nemo/core/utils/cuda_python_utils.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import contextlib
+import inspect
 
 import numpy as np
 import torch
@@ -125,7 +126,8 @@ def with_conditional_node(while_loop_kernel, while_loop_args, while_loop_conditi
     from cuda.bindings import driver as cuda
     from cuda.bindings import runtime as cudart
 
-    capture_status, _, graph, _, _, _ = cu_call(
+    # 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
@@ -144,7 +146,8 @@ def with_conditional_node(while_loop_kernel, while_loop_args, while_loop_conditi
         0,
     )
 
-    capture_status, _, graph, dependencies, _, _ = cu_call(
+    # 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
@@ -168,18 +171,41 @@ def with_conditional_node(while_loop_kernel, while_loop_args, while_loop_conditi
     # 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))
-    (node,) = cu_call(cuda.cuGraphAddNode(graph, dependencies, None, len(dependencies), driver_params))
+    # depending on cuda-python version, number of parameters vary
+    num_cuda_graph_add_node_params = len(inspect.signature(cuda.cuGraphAddNode).parameters)
+    if num_cuda_graph_add_node_params == 5:
+        (node,) = cu_call(cuda.cuGraphAddNode(graph, dependencies, None, len(dependencies), driver_params))
+    elif num_cuda_graph_add_node_params == 4:
+        (node,) = cu_call(cuda.cuGraphAddNode(graph, dependencies, len(dependencies), driver_params))
+    else:
+        raise NeMoCUDAPythonException("Unexpected number of parameters for `cuGraphAddNode`")
     body_graph = driver_params.conditional.phGraph_out[0]
 
-    cu_call(
-        cudart.cudaStreamUpdateCaptureDependencies(
-            torch.cuda.current_stream(device=device).cuda_stream,
-            [node],
-            None,
-            1,
-            cudart.cudaStreamUpdateCaptureDependenciesFlags.cudaStreamSetCaptureDependencies,
-        )
+    # depending on cuda-python version, number of parameters vary
+    num_cuda_stream_update_capture_dependencies_params = len(
+        inspect.signature(cudart.cudaStreamUpdateCaptureDependencies).parameters
     )
+    if num_cuda_stream_update_capture_dependencies_params == 5:
+        cu_call(
+            cudart.cudaStreamUpdateCaptureDependencies(
+                torch.cuda.current_stream(device=device).cuda_stream,
+                [node],
+                None,
+                1,
+                cudart.cudaStreamUpdateCaptureDependenciesFlags.cudaStreamSetCaptureDependencies,
+            )
+        )
+    elif num_cuda_stream_update_capture_dependencies_params == 4:
+        cu_call(
+            cudart.cudaStreamUpdateCaptureDependencies(
+                torch.cuda.current_stream(device=device).cuda_stream,
+                [node],
+                1,
+                cudart.cudaStreamUpdateCaptureDependenciesFlags.cudaStreamSetCaptureDependencies,
+            )
+        )
+    else:
+        raise NeMoCUDAPythonException("Unexpected number of parameters for `cudaStreamUpdateCaptureDependencies`")
     body_stream = torch.cuda.Stream(device)
     previous_stream = torch.cuda.current_stream(device=device)
     cu_call(

Original file line number	Diff line number	Diff line change
`@@ -586,7 +586,8 @@ def _full_graph_compile(self):`
`586`	`586`	`torch.cuda.graph(self.full_graph, stream=stream_for_graph, capture_error_mode="thread_local"),`
`587`	`587`	`):`
`588`	`588`	`self._before_process_batch()`
`589`		`- capture_status, _, graph, _, _, _ = cu_call(`
	`589`	`+ # NB: depending on cuda-python version, cudaStreamGetCaptureInfo can return either 5 or 6 elements`
	`590`	`+ capture_status, _, graph, *_ = cu_call(`
`590`	`591`	`cudart.cudaStreamGetCaptureInfo(torch.cuda.current_stream(device=self.state.device).cuda_stream)`
`591`	`592`	`)`
`592`	`593`
Original file line number	Diff line number	Diff line change
`@@ -874,7 +874,8 @@ def _full_graph_compile(self):`
`874`	`874`	`):`
`875`	`875`	`self._before_loop()`
`876`	`876`
`877`		`- capture_status, _, graph, _, _, _ = cu_call(`
	`877`	`+ # NB: depending on cuda-python version, cudaStreamGetCaptureInfo can return either 5 or 6 elements`
	`878`	`+ capture_status, _, graph, *_ = cu_call(`
`878`	`879`	`cudart.cudaStreamGetCaptureInfo(torch.cuda.current_stream(device=self.state.device).cuda_stream)`
`879`	`880`	`)`
`880`	`881`	`assert capture_status == cudart.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive`
Original file line number	Diff line number	Diff line change
`@@ -205,7 +205,8 @@ def _reinitialize(self, max_time, batch_size, encoder_output, encoder_output_len`
`205`	`205`	`# Get max sequence length`
`206`	`206`	`self.max_out_len_t = self.encoder_output_length.max()`
`207`	`207`
`208`		`- capture_status, _, graph, _, _, _ = cu_call(`
	`208`	`+ # NB: depending on cuda-python version, cudaStreamGetCaptureInfo can return either 5 or 6 elements`
	`209`	`+ capture_status, _, graph, *_ = cu_call(`
`209`	`210`	`cudart.cudaStreamGetCaptureInfo(torch.cuda.current_stream(device=self.device).cuda_stream)`
`210`	`211`	`)`
`211`	`212`	`assert capture_status == cudart.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive`
Original file line number	Diff line number	Diff line change
`@@ -906,7 +906,8 @@ def _full_graph_compile(self):`
`906`	`906`	`torch.cuda.graph(self.full_graph, stream=stream_for_graph, capture_error_mode="thread_local"),`
`907`	`907`	`):`
`908`	`908`	`self._before_loop()`
`909`		`- capture_status, _, graph, _, _, _ = cu_call(`
	`909`	`+ # NB: depending on cuda-python version, cudaStreamGetCaptureInfo can return either 5 or 6 elements`
	`910`	`+ capture_status, _, graph, *_ = cu_call(`
`910`	`911`	`cudart.cudaStreamGetCaptureInfo(torch.cuda.current_stream(device=self.state.device).cuda_stream)`
`911`	`912`	`)`
`912`	`913`
Original file line number	Diff line number	Diff line change
`@@ -987,7 +987,8 @@ def _full_graph_compile(self):`
`987`	`987`	`torch.cuda.graph(self.full_graph, stream=stream_for_graph, capture_error_mode="thread_local"),`
`988`	`988`	`):`
`989`	`989`	`self._before_loop()`
`990`		`- capture_status, _, graph, _, _, _ = cu_call(`
	`990`	`+ # NB: depending on cuda-python version, cudaStreamGetCaptureInfo can return either 5 or 6 elements`
	`991`	`+ capture_status, _, graph, *_ = cu_call(`
`991`	`992`	`cudart.cudaStreamGetCaptureInfo(torch.cuda.current_stream(device=self.state.device).cuda_stream)`
`992`	`993`	`)`
`993`	`994`
Original file line number	Diff line number	Diff line change
`@@ -968,7 +968,8 @@ def _full_graph_compile(self):`
`968`	`968`	`):`
`969`	`969`	`self._before_outer_loop()`
`970`	`970`
`971`		`- capture_status, _, graph, _, _, _ = cu_call(`
	`971`	`+ # NB: depending on cuda-python version, cudaStreamGetCaptureInfo can return either 5 or 6 elements`
	`972`	`+ capture_status, _, graph, *_ = cu_call(`
`972`	`973`	`cudart.cudaStreamGetCaptureInfo(torch.cuda.current_stream(device=self.state.device).cuda_stream)`
`973`	`974`	`)`
`974`	`975`	`assert capture_status == cudart.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive`
Original file line number	Diff line number	Diff line change
`@@ -1041,7 +1041,8 @@ def _full_graph_compile(self):`
`1041`	`1041`	`):`
`1042`	`1042`	`self._before_outer_loop()`
`1043`	`1043`
`1044`		`- capture_status, _, graph, _, _, _ = cu_call(`
	`1044`	`+ # NB: depending on cuda-python version, cudaStreamGetCaptureInfo can return either 5 or 6 elements`
	`1045`	`+ capture_status, _, graph, *_ = cu_call(`
`1045`	`1046`	`cudart.cudaStreamGetCaptureInfo(torch.cuda.current_stream(device=self.state.device).cuda_stream)`
`1046`	`1047`	`)`
`1047`	`1048`	`assert capture_status == cudart.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive`