From 8d80dab74d879d57bd1d325a1c0d305d8a482b8d Mon Sep 17 00:00:00 2001 From: MahmoudAshraf97 Date: Thu, 30 Jul 2026 10:25:32 -0400 Subject: [PATCH 1/6] Clamp the RNN-T gradient in its FP32 register The three Numba grad kernels wrote `grad` to the output buffer, read it straight back, clamped it, and wrote it again. Clamp the value while it is still an FP32 register instead: this drops a global read plus a redundant write per vocabulary element per (b, t, u), and stops narrow output dtypes from being rounded twice. FP32 loss and gradients are bit-identical before and after. Signed-off-by: MahmoudAshraf97 --- .../utils/cuda_utils/gpu_rnnt_kernel.py | 39 +++++++++---------- 1 file changed, 18 insertions(+), 21 deletions(-) diff --git a/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt_kernel.py b/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt_kernel.py index 219e9d0453b2..87eb0c6beeba 100644 --- a/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt_kernel.py +++ b/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt_kernel.py @@ -392,16 +392,15 @@ def compute_grad_kernel( # multiplying (1.0 + fastemit_lambda) with result. grad -= math.exp(math.log1p(fastemit_lambda) + alphas[col] + logpk - logll[mb] + betas[col + 1]) + # clamp gradient (if needed) while it is still an FP32 register, so that + # narrow `grads` dtypes are not rounded twice. + if clamp > 0.0: + grad = min(grad, clamp) + grad = max(grad, -clamp) + # update grads[b, t, u, v] = grad grads[col * alphabet_size + idx] = grad - # clamp gradient (if needed) - if clamp > 0.0: - g = grads[col * alphabet_size + idx] - g = min(g, clamp) - g = max(g, -clamp) - grads[col * alphabet_size + idx] = g - # update internal index through the thread_buffer; # until idx < V + 1, such that entire vocabulary has been updated. idx += GPU_RNNT_THREAD_SIZE @@ -870,16 +869,15 @@ def compute_multiblank_grad_kernel( math.log1p(fastemit_lambda) + alphas[col] + logpk - sigma - logll[mb] + betas[col + 1] ) + # clamp gradient (if needed) while it is still an FP32 register, so that + # narrow `grads` dtypes are not rounded twice. + if clamp > 0.0: + grad = min(grad, clamp) + grad = max(grad, -clamp) + # update grads[b, t, u, v] = grad grads[col * alphabet_size + idx] = grad - # clamp gradient (if needed) - if clamp > 0.0: - g = grads[col * alphabet_size + idx] - g = min(g, clamp) - g = max(g, -clamp) - grads[col * alphabet_size + idx] = g - # update internal index through the thread_buffer; # until idx < V + 1, such that entire vocabulary has been updated. idx += GPU_RNNT_THREAD_SIZE @@ -1424,16 +1422,15 @@ def compute_tdt_grad_kernel( + duration_acts[col * num_durations + i] ) + # clamp gradient (if needed) while it is still an FP32 register, so that + # narrow `label_grads` dtypes are not rounded twice. + if clamp > 0.0: + grad = min(grad, clamp) + grad = max(grad, -clamp) + # update grads[b, t, u, v] = grad label_grads[col * alphabet_size + idx] = grad - # clamp gradient (if needed) - if clamp > 0.0: - g = label_grads[col * alphabet_size + idx] - g = min(g, clamp) - g = max(g, -clamp) - label_grads[col * alphabet_size + idx] = g - # update internal index through the thread_buffer; # until idx < V + 1, such that entire vocabulary has been updated. idx += GPU_RNNT_THREAD_SIZE From ef630b6a87db0d153cbc1ef92ad3f1e8685a81fe Mon Sep 17 00:00:00 2001 From: MahmoudAshraf97 Date: Thu, 30 Jul 2026 10:25:32 -0400 Subject: [PATCH 2/6] Promote Numba RNN-T activation reads to FP32 Numba-CUDA cannot implicitly unify BF16 with the FP32 dynamic-programming state, so BF16 activations make kernel type inference diverge and the launch fails with RecursionError. Widen activation reads at the point of use in logp() and logp_duration(), which keeps `acts` narrow in memory and does the arithmetic in FP32. Also allocate the CTA reduction scratch as the module FP32 dtype rather than the activation dtype, so reductions no longer accumulate in the input precision. Standard RNN-T now runs under BF16. Multi-blank and TDT still fail earlier, in the workspace allocation. FP32 results are unchanged. Signed-off-by: MahmoudAshraf97 --- .../utils/cuda_utils/gpu_rnnt_kernel.py | 5 +++-- .../numba/rnnt_loss/utils/cuda_utils/reduce.py | 16 ++++++++++------ 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt_kernel.py b/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt_kernel.py index 87eb0c6beeba..bf9bf77c6d33 100644 --- a/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt_kernel.py +++ b/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt_kernel.py @@ -28,6 +28,7 @@ import math +import numba import torch from numba import cuda @@ -61,13 +62,13 @@ def logp( The sum of logprobs[mb, t, u, v] + denom[mb, t, u] """ col = (mb * maxT + t) * maxU + u - return denom[col] + acts[col * alphabet_size + v] + return denom[col] + numba.float32(acts[col * alphabet_size + v]) @cuda.jit(device=True, inline=True) def logp_duration(acts: torch.Tensor, maxT: int, maxU: int, num_durations: int, mb: int, t: int, u: int, v: int): col = (mb * maxT + t) * maxU + u - return acts[col * num_durations + v] + return numba.float32(acts[col * num_durations + v]) @cuda.jit() diff --git a/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/reduce.py b/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/reduce.py index c72026ae2fee..b98add86fc5f 100644 --- a/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/reduce.py +++ b/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/reduce.py @@ -29,6 +29,7 @@ import enum import math +import numba import torch from numba import cuda @@ -37,6 +38,9 @@ warp_size = global_constants.warp_size() dtype = global_constants.dtype() +# RNN-T dynamic-programming state is FP32. Promote narrow activation reads +# explicitly because Numba-CUDA cannot implicitly unify BF16 with FP32. + CTA_REDUCE_SIZE = 128 @@ -147,13 +151,13 @@ def _reduce_rows(I_opid: int, R_opid: int, acts, output, num_rows: int): col = cuda.blockIdx.x # allocate shared thread memory - storage = cuda.shared.array(shape=(CTA_REDUCE_SIZE,), dtype=acts.dtype) + storage = cuda.shared.array(shape=(CTA_REDUCE_SIZE,), dtype=dtype) max = output[col] # // Each block works on a column if idx < num_rows: - curr = acts[col * num_rows + idx] - max + curr = numba.float32(acts[col * num_rows + idx]) - max if I_opid == 0: curr = rnnt_helper.exponential(curr) else: @@ -162,7 +166,7 @@ def _reduce_rows(I_opid: int, R_opid: int, acts, output, num_rows: int): idx += CTA_REDUCE_SIZE while idx < num_rows: - activation_ = acts[col * num_rows + idx] - max + activation_ = numba.float32(acts[col * num_rows + idx]) - max if I_opid == 0 and R_opid == 0: curr = rnnt_helper.add(curr, rnnt_helper.exponential(activation_)) @@ -212,13 +216,13 @@ def _reduce_minus(I_opid: int, R_opid: int, acts, output, num_rows: int): col = cuda.blockIdx.x # allocate shared thread memory - storage = cuda.shared.array(shape=(CTA_REDUCE_SIZE,), dtype=acts.dtype) + storage = cuda.shared.array(shape=(CTA_REDUCE_SIZE,), dtype=dtype) max = output[col] # // Each block works on a column if idx < num_rows: - curr = acts[col * num_rows + idx] - max + curr = numba.float32(acts[col * num_rows + idx]) - max if I_opid == 0: curr = rnnt_helper.exponential(curr) else: @@ -227,7 +231,7 @@ def _reduce_minus(I_opid: int, R_opid: int, acts, output, num_rows: int): idx += CTA_REDUCE_SIZE while idx < num_rows: - activation_ = acts[col * num_rows + idx] - max + activation_ = numba.float32(acts[col * num_rows + idx]) - max if I_opid == 0 and R_opid == 0: curr = rnnt_helper.add(curr, rnnt_helper.exponential(activation_)) From ecf10114ad5105529f9d1fb2c17363927449143b Mon Sep 17 00:00:00 2001 From: MahmoudAshraf97 Date: Thu, 30 Jul 2026 10:25:32 -0400 Subject: [PATCH 3/6] Route Numba grad kernels through the logp helpers compute_grad_kernel, compute_multiblank_grad_kernel and compute_tdt_grad_kernel open-coded `denom[col] + acts[...]` and `duration_acts[...]` instead of calling logp() and logp_duration(). Use the helpers, which removes the duplicated index arithmetic and picks up their FP32 promotion. No uncast activation read remains in the Numba RNN-T kernels. FP32 results are unchanged. Signed-off-by: MahmoudAshraf97 --- .../utils/cuda_utils/gpu_rnnt_kernel.py | 32 +++++++++++-------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt_kernel.py b/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt_kernel.py index bf9bf77c6d33..a7ffa57de7f9 100644 --- a/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt_kernel.py +++ b/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt_kernel.py @@ -352,7 +352,7 @@ def compute_grad_kernel( while idx < alphabet_size: # remember, `col` represents the tri-index [b, t, u] # therefore; logpk = denom[b, t, u] + acts[b, t, u, v] - logpk = denom[col] + acts[col * alphabet_size + idx] + logpk = logp(denom, acts, maxT, maxU, alphabet_size, mb, t, u, idx) # initialize the grad of the sample acts[b, t, u, v] grad = math.exp(alphas[col] + betas[col] + logpk - logll[mb]) @@ -364,7 +364,7 @@ def compute_grad_kernel( if fastemit_lambda > 0.0 and u < U - 1: fastemit_grad = fastemit_lambda * math.exp( alphas[col] # alphas(t, u) - + (denom[col] + acts[col * alphabet_size + labels[u]]) # y_hat(t, u) + + logp(denom, acts, maxT, maxU, alphabet_size, mb, t, u, labels[u]) # y_hat(t, u) + betas[col + 1] # betas(t, u+1) + logpk # log Pr(k|t, u) - logll[mb] # total log likelihood for normalization @@ -804,7 +804,7 @@ def compute_multiblank_grad_kernel( while idx < alphabet_size: # remember, `col` represents the tri-index [b, t, u] # therefore; logpk = denom[b, t, u] + acts[b, t, u, v] - logpk = denom[col] + acts[col * alphabet_size + idx] + logpk = logp(denom, acts, maxT, maxU, alphabet_size, mb, t, u, idx) # initialize the grad of the sample acts[b, t, u, v] grad = math.exp(alphas[col] + betas[col] + logpk - logll[mb]) @@ -820,7 +820,7 @@ def compute_multiblank_grad_kernel( if fastemit_lambda > 0.0 and u < U - 1: fastemit_grad = fastemit_lambda * math.exp( alphas[col] # alphas(t, u) - + (denom[col] + acts[col * alphabet_size + labels[u]]) + + logp(denom, acts, maxT, maxU, alphabet_size, mb, t, u, labels[u]) # y_hat(t, u) + betas[col + 1] # betas(t, u+1) + logpk # log Pr(k|t, u) - sigma @@ -1321,13 +1321,13 @@ def compute_tdt_grad_kernel( if t < T and u < U: logpk_blank = ( - denom[col] + acts[col * alphabet_size + blank_] - sigma + logp(denom, acts, maxT, maxU, alphabet_size, mb, t, u, blank_) - sigma ) # whenever sigma is used, it is for logit under-normalization. if idx < num_durations: grad = 0.0 if t + durations[idx] < T and u < U - 1: # for label - logpk_label = denom[col] + acts[col * alphabet_size + labels[u]] - sigma + logpk_label = logp(denom, acts, maxT, maxU, alphabet_size, mb, t, u, labels[u]) - sigma grad -= math.exp(alphas[col] + betas[col + 1 + durations[idx] * maxU] + logpk_label - logll[mb]) if t + durations[idx] < T and durations[idx] > 0: # for blank in the middle @@ -1336,7 +1336,7 @@ def compute_tdt_grad_kernel( if t + durations[idx] == T and u == U - 1 and durations[idx] > 0: # for blank as the last symbol grad -= math.exp(alphas[col] + logpk_blank - logll[mb]) - grad = grad * math.exp(duration_acts[col * num_durations + idx]) + grad = grad * math.exp(logp_duration(duration_acts, maxT, maxU, num_durations, mb, t, u, idx)) duration_grads[col * num_durations + idx] = grad # For cuda kernels, maximum number of threads per block is limited to some value. @@ -1349,7 +1349,7 @@ def compute_tdt_grad_kernel( while idx < alphabet_size: # remember, `col` represents the tri-index [b, t, u] # therefore; logpk = denom[b, t, u] + acts[b, t, u, v] - logpk = denom[col] + acts[col * alphabet_size + idx] + logpk = logp(denom, acts, maxT, maxU, alphabet_size, mb, t, u, idx) # initialize the grad of the sample acts[b, t, u, v] grad = math.exp(alphas[col] + betas[col] + logpk - logll[mb]) @@ -1365,8 +1365,10 @@ def compute_tdt_grad_kernel( if t + durations[i] < T: fastemit_grad += fastemit_lambda * math.exp( alphas[col] # alphas(t, u) - + (denom[col] + acts[col * alphabet_size + labels[u]]) # log prob of token emission - + duration_acts[col * num_durations + i] # duration log-prob + + logp( + denom, acts, maxT, maxU, alphabet_size, mb, t, u, labels[u] + ) # log prob of token emission + + logp_duration(duration_acts, maxT, maxU, num_durations, mb, t, u, i) # duration log-prob + betas[col + 1 + durations[i] * maxU] # betas(t, u+1) + logpk # log Pr(k|t, u) - sigma # for logit under-normalization @@ -1386,7 +1388,11 @@ def compute_tdt_grad_kernel( continue if t == T - durations[i]: grad -= math.exp( - alphas[col] + logpk - sigma - logll[mb] + duration_acts[col * num_durations + i] + alphas[col] + + logpk + - sigma + - logll[mb] + + logp_duration(duration_acts, maxT, maxU, num_durations, mb, t, u, i) ) # grad of blank across t < T; @@ -1402,7 +1408,7 @@ def compute_tdt_grad_kernel( - sigma - logll[mb] + betas[col + maxU * durations[i]] - + duration_acts[col * num_durations + i] + + logp_duration(duration_acts, maxT, maxU, num_durations, mb, t, u, i) ) # grad of correct token across u < U; @@ -1420,7 +1426,7 @@ def compute_tdt_grad_kernel( - sigma - logll[mb] + betas[col + 1 + maxU * durations[i]] - + duration_acts[col * num_durations + i] + + logp_duration(duration_acts, maxT, maxU, num_durations, mb, t, u, i) ) # clamp gradient (if needed) while it is still an FP32 register, so that From 35cb7a70bb27b9c669d17b54a3249a3c5f9f5864 Mon Sep 17 00:00:00 2001 From: MahmoudAshraf97 Date: Thu, 30 Jul 2026 10:25:32 -0400 Subject: [PATCH 4/6] Allocate the multi-blank and TDT GPU workspace in FP32 The workspace backs FP32 dynamic-programming state -- softmax denominator, alphas, betas, log-likelihoods -- but multiblank_rnnt_loss_gpu and tdt_loss_gpu sized it with the activation dtype. With BF16 activations the denominator buffer reached the reduction kernel as a 2-byte opaque dtype and the launch failed with NumbaNotImplementedError. rnnt_loss_gpu already hardcodes FP32; make the other two consistent. get_workspace_size returns an element count, so this changes only the buffer dtype, not its length. Multi-blank and TDT now run under BF16. FP32 results are unchanged. Signed-off-by: MahmoudAshraf97 --- nemo/collections/asr/parts/numba/rnnt_loss/rnnt.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/nemo/collections/asr/parts/numba/rnnt_loss/rnnt.py b/nemo/collections/asr/parts/numba/rnnt_loss/rnnt.py index 046aea425e20..7d12b705e57f 100644 --- a/nemo/collections/asr/parts/numba/rnnt_loss/rnnt.py +++ b/nemo/collections/asr/parts/numba/rnnt_loss/rnnt.py @@ -299,7 +299,9 @@ def tdt_loss_gpu( # Select GPU index cuda.select_device(label_acts.device.index) - gpu_workspace = torch.zeros(gpu_size, device=label_acts.device, dtype=label_acts.dtype, requires_grad=False) + # The workspace holds FP32 dynamic-programming state (denominator, alphas, betas, + # log-likelihoods), so it must not inherit a narrow activation dtype. + gpu_workspace = torch.zeros(gpu_size, device=label_acts.device, dtype=torch.float32, requires_grad=False) tdt_workspace = torch.zeros(len(durations), device=label_acts.device, dtype=torch.long, requires_grad=False) @@ -423,7 +425,9 @@ def multiblank_rnnt_loss_gpu( # Select GPU index cuda.select_device(acts.device.index) - gpu_workspace = torch.zeros(gpu_size, device=acts.device, dtype=acts.dtype, requires_grad=False) + # The workspace holds FP32 dynamic-programming state (denominator, alphas, betas, + # log-likelihoods), so it must not inherit a narrow activation dtype. + gpu_workspace = torch.zeros(gpu_size, device=acts.device, dtype=torch.float32, requires_grad=False) big_blank_workspace = torch.zeros( len(big_blank_durations), device=acts.device, dtype=torch.long, requires_grad=False From 6933bfb847660b76c181cbb6ad419be4ea424bb3 Mon Sep 17 00:00:00 2001 From: MahmoudAshraf97 Date: Thu, 30 Jul 2026 10:25:32 -0400 Subject: [PATCH 5/6] Test Numba RNN-T losses with BF16 activations Covers standard, multi-blank and TDT across fastemit and clamp settings, asserting each compiles under BF16 and agrees with FP32. The regression these guard is a kernel compilation failure rather than numerical drift, so the assertion that matters most is that the launch succeeds at all. Gated only on Numba CUDA support: BF16 needs no NUMBA_CUDA_USE_NVIDIA_BINDING at the Numba level -- that variable gates whether NeMo routes narrow activations to the loss, not whether the kernels can consume them -- so these run in ordinary CI rather than silently skipping. Signed-off-by: MahmoudAshraf97 --- .../asr/numba/rnnt_loss/test_rnnt_pytorch.py | 76 +++++++++++++++++++ 1 file changed, 76 insertions(+) diff --git a/tests/collections/asr/numba/rnnt_loss/test_rnnt_pytorch.py b/tests/collections/asr/numba/rnnt_loss/test_rnnt_pytorch.py index 9f38bf6dbe8a..a2078eef08d7 100644 --- a/tests/collections/asr/numba/rnnt_loss/test_rnnt_pytorch.py +++ b/tests/collections/asr/numba/rnnt_loss/test_rnnt_pytorch.py @@ -625,5 +625,81 @@ def test_case_fixed_case_act_label(self, device): assert np.allclose(pt_grads, expected_grads, rtol=1e-2), "td gradient mismatch." +class TestRNNTLossNumbaBFloat16: + """BF16 activations must compile and agree with FP32. + + Numba-CUDA cannot implicitly unify BF16 with the FP32 dynamic-programming state, so a + regression here surfaces as a kernel compilation failure (``RecursionError`` from type + inference, or ``NumbaNotImplementedError`` when a BF16 buffer reaches a kernel argument) + rather than a numerical mismatch. + """ + + @staticmethod + def _inputs(dtype, vocab, batch=3, max_source=7, max_target=4): + torch.manual_seed(3) + acts = torch.randn(batch, max_source, max_target + 1, vocab, device='cuda', dtype=dtype, requires_grad=True) + labels = torch.randint(0, vocab - 4, (batch, max_target), device='cuda', dtype=torch.int64) + source_lengths = torch.full((batch,), max_source, device='cuda', dtype=torch.int64) + target_lengths = torch.full((batch,), max_target, device='cuda', dtype=torch.int64) + return acts, labels, source_lengths, target_lengths + + def _compare_against_float32(self, build_loss, vocab): + costs, grads = {}, {} + for dtype in (torch.float32, torch.bfloat16): + acts, labels, source_lengths, target_lengths = self._inputs(dtype, vocab) + cost = build_loss()(acts, labels, source_lengths, target_lengths) + grads[dtype] = torch.autograd.grad(cost.sum(), acts)[0].float() + costs[dtype] = cost.float() + + assert torch.isfinite(grads[torch.bfloat16]).all() + assert torch.count_nonzero(grads[torch.bfloat16]) + torch.testing.assert_close(costs[torch.bfloat16], costs[torch.float32], atol=2e-2, rtol=2e-2) + torch.testing.assert_close(grads[torch.bfloat16], grads[torch.float32], atol=2e-2, rtol=2e-2) + + @pytest.mark.unit + @pytest.mark.parametrize('fastemit_lambda', [0.0, 0.01]) + @pytest.mark.parametrize('clamp', [-1.0, 0.02]) + def test_rnnt_bfloat16_matches_float32(self, fastemit_lambda, clamp): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + self._compare_against_float32( + lambda: RNNTLossNumba(blank=7, reduction='none', fastemit_lambda=fastemit_lambda, clamp=clamp), + vocab=8, + ) + + @pytest.mark.unit + @pytest.mark.parametrize('fastemit_lambda', [0.0, 0.01]) + @pytest.mark.parametrize('clamp', [-1.0, 0.02]) + def test_multiblank_rnnt_bfloat16_matches_float32(self, fastemit_lambda, clamp): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + self._compare_against_float32( + lambda: MultiblankRNNTLossNumba( + blank=9, + big_blank_durations=[2, 4], + reduction='none', + fastemit_lambda=fastemit_lambda, + clamp=clamp, + sigma=0.05, + ), + vocab=10, + ) + + @pytest.mark.unit + @pytest.mark.parametrize('fastemit_lambda', [0.0, 0.01]) + @pytest.mark.parametrize('clamp', [-1.0, 0.02]) + def test_tdt_bfloat16_matches_float32(self, fastemit_lambda, clamp): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + self._compare_against_float32( + lambda: TDTLossNumba( + blank=7, + durations=[0, 1, 2], + reduction='none', + fastemit_lambda=fastemit_lambda, + clamp=clamp, + sigma=0.05, + ), + vocab=11, + ) + + if __name__ == "__main__": pytest.main([__file__]) From 74ecfa74745551799f3577b002b7c2d83c296612 Mon Sep 17 00:00:00 2001 From: MahmoudAshraf97 Date: Thu, 30 Jul 2026 10:28:07 -0400 Subject: [PATCH 6/6] Use single-hash block comments in the Numba RNN-T wrapper The `### ... ###` banners trip flake8's E266. Linting runs over the files a change touches, so these pre-existing violations otherwise fail CI for any branch that edits this module. Comment text is unchanged. Signed-off-by: MahmoudAshraf97 --- .../asr/parts/numba/rnnt_loss/rnnt.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/nemo/collections/asr/parts/numba/rnnt_loss/rnnt.py b/nemo/collections/asr/parts/numba/rnnt_loss/rnnt.py index 7d12b705e57f..384b67b7f04c 100644 --- a/nemo/collections/asr/parts/numba/rnnt_loss/rnnt.py +++ b/nemo/collections/asr/parts/numba/rnnt_loss/rnnt.py @@ -86,7 +86,7 @@ def rnnt_loss_cpu( cpu_workspace = torch.zeros(gpu_size, device=log_probs.device, dtype=log_probs.dtype, requires_grad=False) - ### VIEW TENSORS AS VECTORS FOR POINTER INDEXING ### + # VIEW TENSORS AS VECTORS FOR POINTER INDEXING log_probs, acts_shape = rnnt_helper.flatten_tensor(log_probs) flat_labels, labels_shape = rnnt_helper.flatten_tensor(flat_labels) @@ -116,7 +116,7 @@ def rnnt_loss_cpu( raise RuntimeError("Could not calculate forward scores") else: - ### FLATTEN GRAD TENSOR ### + # FLATTEN GRAD TENSOR grads, grads_shape = rnnt_helper.flatten_tensor(grads) status = wrapper.cost_and_grad( @@ -188,7 +188,7 @@ def rnnt_loss_gpu( cuda.select_device(acts.device.index) gpu_workspace = torch.zeros(gpu_size, device=acts.device, dtype=torch.float32, requires_grad=False) - ### VIEW TENSORS AS VECTORS FOR POINTER INDEXING ### + # VIEW TENSORS AS VECTORS FOR POINTER INDEXING acts, acts_shape = rnnt_helper.flatten_tensor(acts) wrapper = gpu_rnnt.GPURNNT( @@ -217,7 +217,7 @@ def rnnt_loss_gpu( raise RuntimeError("Could not calculate forward scores") else: - ### FLATTEN GRAD TENSOR ### + # FLATTEN GRAD TENSOR grads, grads_shape = rnnt_helper.flatten_tensor(grads) status = wrapper.cost_and_grad( @@ -308,7 +308,7 @@ def tdt_loss_gpu( for i in range(0, len(durations)): tdt_workspace[i] = durations[i] - ### VIEW TENSORS AS VECTORS FOR POINTER INDEXING ### + # VIEW TENSORS AS VECTORS FOR POINTER INDEXING label_acts, label_acts_shape = rnnt_helper.flatten_tensor(label_acts) duration_acts, duration_acts_shape = rnnt_helper.flatten_tensor(duration_acts) @@ -343,7 +343,7 @@ def tdt_loss_gpu( raise RuntimeError("Could not calculate forward scores") else: - ### FLATTEN GRAD TENSOR ### + # FLATTEN GRAD TENSOR label_grads, label_grads_shape = rnnt_helper.flatten_tensor(label_grads) duration_grads, duration_grads_shape = rnnt_helper.flatten_tensor(duration_grads) @@ -436,7 +436,7 @@ def multiblank_rnnt_loss_gpu( for i in range(0, len(big_blank_durations)): big_blank_workspace[i] = big_blank_durations[i] - ### VIEW TENSORS AS VECTORS FOR POINTER INDEXING ### + # VIEW TENSORS AS VECTORS FOR POINTER INDEXING acts, acts_shape = rnnt_helper.flatten_tensor(acts) wrapper = gpu_rnnt.MultiblankGPURNNT( @@ -468,7 +468,7 @@ def multiblank_rnnt_loss_gpu( raise RuntimeError("Could not calculate forward scores") else: - ### FLATTEN GRAD TENSOR ### + # FLATTEN GRAD TENSOR grads, grads_shape = rnnt_helper.flatten_tensor(grads) status = wrapper.cost_and_grad(