Support BF16 activations in the Numba RNN-T losses

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(
@@ -299,14 +299,16 @@ 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)
 
     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)
 
@@ -341,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)
 
@@ -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
@@ -432,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(
@@ -464,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(

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()
@@ -351,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])
 
@@ -363,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
@@ -392,16 +393,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
@@ -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
@@ -870,16 +870,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
@@ -1322,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
@@ -1337,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.
@@ -1350,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])
 
@@ -1366,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
@@ -1387,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;
@@ -1403,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;
@@ -1421,19 +1426,18 @@ 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
+            # 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

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_))