Optimized Graph-Transducer Implementation

nemo/collections/asr/parts/k2/graph_transducer.py

@@ -13,14 +13,20 @@
 # limitations under the License.
 
 import abc
+import logging
 from contextlib import nullcontext
 from typing import ContextManager
+
 import torch
 import torch.nn.functional as F
 
+from nemo.collections.common.optional_libs import TRITON_AVAILABLE
 from nemo.core.classes.loss import Loss
 from nemo.core.utils.k2_guard import k2
 
+if TRITON_AVAILABLE:
+    from nemo.collections.asr.parts.k2.rnnt_logprobs_triton import rnnt_logprobs_triton
+
 
 def force_float32_context() -> ContextManager:
     """Get context manager to force float32 precision in autocast mode."""
@@ -146,11 +152,13 @@ def get_graphs_batched(
         batch_size = logits_lengths.shape[0]
         with torch.no_grad():
             if self.use_grid_implementation:
+                logits_lengths_list = logits_lengths.tolist()
+                target_lengths_list = target_lengths.tolist()
                 return k2.create_fsa_vec(
                     [
                         self.get_grid(
-                            units_tensor=targets[i, : target_lengths[i].item()],
-                            num_frames=logits_lengths[i].item(),
+                            units_tensor=targets[i, : target_lengths_list[i]],
+                            num_frames=logits_lengths_list[i],
                             vocab_size=vocab_size,
                         )
                         for i in range(batch_size)
@@ -175,22 +183,20 @@ def get_graphs_batched(
                 k2.create_fsa_vec(text_fsas), k2.create_fsa_vec(temporal_fsas), treat_epsilons_specially=False
             )
             if self.connect_composed:
-                k2.connect(target_fsas_vec)
+                target_fsas_vec = k2.connect(target_fsas_vec)
         return target_fsas_vec
 
-    def get_logits_indices(self, target_fsas_vec: k2.Fsa, logits_shape: torch.Size) -> torch.Tensor:
+    def get_batch_indices(self, target_fsas_vec: k2.Fsa) -> torch.Tensor:
         """
         Get indices of flatten logits for each arc in the lattices.
 
         Args:
             target_fsas_vec: batch of target FSAs with lattices
-            logits_shape: shape of the logits tensor
 
         Returns:
             1d tensor with indices
         """
-        # logits_shape: B x Time x Text+1 x Labels
-        batch_size = logits_shape[0]
+        batch_size = target_fsas_vec.shape[0]
         device = target_fsas_vec.device
         scores_to_batch_i = torch.repeat_interleave(
             torch.arange(batch_size, device=device, dtype=torch.int64),
@@ -199,6 +205,21 @@ def get_logits_indices(self, target_fsas_vec: k2.Fsa, logits_shape: torch.Size)
                 device=device,
             ),
         )
+        return scores_to_batch_i
+
+    def get_logits_indices(self, target_fsas_vec: k2.Fsa, logits_shape: torch.Size) -> torch.Tensor:
+        """
+        Get indices of flatten logits for each arc in the lattices.
+
+        Args:
+            target_fsas_vec: batch of target FSAs with lattices
+            logits_shape: shape of the logits tensor
+
+        Returns:
+            1d tensor with indices
+        """
+        # logits_shape: B x Time x Text+1 x Labels
+        scores_to_batch_i = self.get_batch_indices(target_fsas_vec=target_fsas_vec)
         indices = (
             scores_to_batch_i * logits_shape[1] * logits_shape[2] * logits_shape[3]  # Batch
             + target_fsas_vec.aux_labels.to(torch.int64) * logits_shape[2] * logits_shape[3]  # Time indices
@@ -222,6 +243,8 @@ def __init__(
         connect_composed=False,
         double_scores=False,
         cast_to_float32=False,
+        return_graph=False,
+        use_triton=True,
     ):
         """
         Init method
@@ -234,6 +257,8 @@ def __init__(
             double_scores: Use calculation of loss in double precision (float64) in the lattice.
                 Does not significantly affect memory usage since the lattice is ~V/2 times smaller than the joint tensor.
             cast_to_float32: Force cast joint tensor to float32 before log-softmax calculation.
+            return_graph: Return graph (along with loss) from `forward` function
+            use_triton: use optimized log probs calculations with Triton (faster and more memory efficient)
         """
         super().__init__(
             use_grid_implementation=use_grid_implementation,
@@ -242,6 +267,10 @@ def __init__(
             cast_to_float32=cast_to_float32,
         )
         self.blank = blank
+        self.return_graph = return_graph
+        self.use_triton = use_triton and TRITON_AVAILABLE
+        if not self.use_triton:
+            logging.warning("Triton is disabled, memory usage can be larger")
 
     def get_unit_schema(self, units_tensor: torch.Tensor, vocab_size: int) -> "k2.Fsa":
         """
@@ -370,13 +399,14 @@ def relabel_states(states: torch.Tensor, n: int, m: int) -> torch.Tensor:
         anti_diag = m + n - 1 - diag
         max_idx = n * m - 1
         cur_diag_idx = i if m > n else m - j - 1
-        states = (
+        new_states = (
             diag.lt(min_mn) * ((diag * (diag + 1) >> 1) + i)
             + torch.logical_and(diag.ge(min_mn), diag.lt(max_mn))
             * ((min_mn * (min_mn + 1) >> 1) + (diag - min_mn) * min_mn + cur_diag_idx)
             + diag.ge(max_mn) * (max_idx - (anti_diag * (anti_diag + 1) >> 1) + m - j)
         )
-        return states
+        torch.where(states >= n * m, states, new_states, out=new_states)
+        return new_states
 
     def get_grid(self, units_tensor: torch.Tensor, num_frames: int, vocab_size: int) -> "k2.Fsa":
         """
@@ -445,13 +475,63 @@ def get_grid(self, units_tensor: torch.Tensor, num_frames: int, vocab_size: int)
         rnnt_graph.unit_positions = unit_positions
         return rnnt_graph
 
+    def get_weighted_graphs(
+        self,
+        logits: torch.Tensor,
+        targets: torch.Tensor,
+        logits_lengths: torch.Tensor,
+        target_lengths: torch.Tensor,
+        use_graph_weight=False,
+    ) -> "k2.Fsa":
+        vocab_size = logits.shape[-1]
+        target_fsas_vec = self.get_graphs_batched(logits_lengths, targets, target_lengths, vocab_size)
+
+        with torch.no_grad():
+            # last transitions in the graph are labeled with -1 label
+            last_transition_mask = target_fsas_vec.labels == -1
+            batch_indices = self.get_batch_indices(target_fsas_vec=target_fsas_vec)
+            time_indices = target_fsas_vec.aux_labels.clone().to(torch.int64)
+            unit_indices = target_fsas_vec.unit_positions.clone().to(torch.int64)
+            text_units = target_fsas_vec.labels.clone().to(torch.int64)
+            # fill in the indices outside the logits with 0, replace later
+            text_units.masked_fill_(last_transition_mask, 0)
+
+        cast_context = force_float32_context() if self.cast_to_float32 else nullcontext()
+        with cast_context:
+            # NB: do not assign scores -> modify, k2 will not update all scores correctly (modify -> assign)
+            if self.use_triton and logits.device.type == "cuda":
+                unit_scores, blank_scores = rnnt_logprobs_triton(
+                    logits=logits,
+                    targets=targets,
+                    blank_id=self.blank,
+                    source_lengths=logits_lengths,
+                    target_lengths=target_lengths,
+                )
+                text_units_blank_mask = text_units == self.blank
+                scores = torch.where(
+                    text_units_blank_mask,
+                    blank_scores[batch_indices, time_indices, unit_indices],
+                    unit_scores[batch_indices, time_indices, unit_indices],
+                ).to(torch.float32)
+                scores[last_transition_mask] = 0.0  # fix weights for the arcs to the last state
+            else:
+                log_probs = F.log_softmax(logits, dim=-1)
+                scores = log_probs[batch_indices, time_indices, unit_indices, text_units].to(torch.float32)
+                scores[last_transition_mask] = 0.0
+
+            if use_graph_weight:
+                target_fsas_vec.scores = target_fsas_vec.scores + scores
+            else:
+                target_fsas_vec.scores = scores
+        return target_fsas_vec
+
     def forward(
         self,
         acts: torch.Tensor,
         labels: torch.Tensor,
         act_lens: torch.Tensor,
         label_lens: torch.Tensor,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, "k2.Fsa"]:
         """
         Compute forward method for RNN-T.
 
@@ -466,26 +546,11 @@ def forward(
         """
         # argument names are consistent with NeMo, see RNNTLoss.forward:
         # self._loss(acts=log_probs, labels=targets, act_lens=input_lengths, label_lens=target_lengths)
-        logits, targets, logits_lengths, target_lengths = acts, labels, act_lens, label_lens
-
-        # logits: B x Time x Text+1 x C
-        vocab_size = logits.shape[-1]
-        target_fsas_vec = self.get_graphs_batched(logits_lengths, targets, target_lengths, vocab_size)
-
-        cast_context = force_float32_context() if self.cast_to_float32 else nullcontext()
-        with cast_context:
-            log_probs = F.log_softmax(logits, dim=-1)
-            with torch.no_grad():
-                indices = self.get_logits_indices(target_fsas_vec, logits.shape)
-                # transition to the last state
-                # use 0 index (for valid index_select) and manually assign score after index_select for this case
-                indices[target_fsas_vec.labels == -1] = 0
-
-            # NB: do not assign scores -> modify, k2 will not update all scores correctly (modify -> assign)
-            scores = log_probs.flatten().index_select(-1, indices)
-            # fix weights for the arcs to the last state
-            scores[target_fsas_vec.labels == -1] = 0
+        target_fsas_vec = self.get_weighted_graphs(
+            logits=acts, targets=labels, logits_lengths=act_lens, target_lengths=label_lens, use_graph_weight=False
+        )
 
-            target_fsas_vec.scores = scores
-            scores = -1 * target_fsas_vec.get_tot_scores(use_double_scores=self.double_scores, log_semiring=True)
-            return scores
+        scores = -1 * target_fsas_vec.get_tot_scores(use_double_scores=self.double_scores, log_semiring=True)
+        if self.return_graph:
+            return scores, target_fsas_vec
+        return scores

nemo/collections/asr/parts/k2/rnnt_logprobs.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2024, 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 torch
+import torch.nn.functional as F
+
+
+def rnnt_logprobs_torch(x: torch.Tensor, targets: torch.Tensor, blank_id: int) -> tuple[torch.Tensor, torch.Tensor]:
+    device = x.device
+    batch_size = x.shape[0]
+    x_log_softmax = F.log_softmax(x, dim=-1)
+    blank_scores = x_log_softmax[..., blank_id]
+    targets = torch.cat((targets, torch.zeros(batch_size, dtype=targets.dtype, device=device).unsqueeze(1)), dim=-1)
+    target_scores = torch.gather(
+        x_log_softmax, dim=-1, index=targets.unsqueeze(1).expand(x.shape[:-1]).unsqueeze(-1)
+    ).squeeze(-1)
+    return target_scores, blank_scores