Replace tensorized bucketization with for loop

Author: tdene

megatron/core/inference/text_generation_controllers/text_generation_controller.py

@@ -6,7 +6,7 @@
 import functools
 import inspect
 from collections import defaultdict
-from typing import Any, Dict, List, Optional, OrderedDict, Tuple, Union
+from typing import Any, Dict, Iterator, List, Optional, OrderedDict, Tuple, Union
 
 import torch
 import torch.nn.functional as F
@@ -115,7 +115,7 @@ def _init_dynamic_sampling_tensors(self):
 
         # Used for inefficient torch sampling.
         if self._sampling_backend == "torch":
-            self._torch_sampling_buckets: List[Tuple] = []
+            self._torch_sampling_buckets: Iterator[Tuple] = []
 
     def tokenize_prompt(self, prompt: str, add_BOS: bool = False) -> List[int]:
         """Utility to tokenize the input prompts.
@@ -613,28 +613,26 @@ def _dynamic_step_sample_bookkeeping(self):
 
         if self._sampling_backend == "torch":
             # Bucketize the core sampling parameters.
-            core_params = torch.stack(
-                (
-                    self._request_metadata["temperature"],
-                    self._request_metadata["top_k"],
-                    self._request_metadata["top_p"],
-                ),
-                dim=1,
-            )
-            _, inv_indices, cnts = torch.unique(
-                core_params, dim=0, return_inverse=True, return_counts=True
-            )
-            order = torch.argsort(inv_indices, stable=True)
-            sampling_buckets = torch.split(order, cnts.tolist())
-            group_reps = torch.stack([indices[0] for indices in sampling_buckets], dim=0)
-            temp_reps = self._request_metadata["temperature"][group_reps].tolist()
-            top_k_reps = self._request_metadata["top_k"][group_reps].tolist()
-            top_p_reps = self._request_metadata["top_p"][group_reps].tolist()
+            # Doing so via list comprehension is orders of magnitude faster than via torch.
+            bucket_map = {}
+
+            # Shorthands for the dictionary comprehension.
+            temp = self._request_metadata["temperature"].tolist()
+            top_k = self._request_metadata["top_k"].tolist()
+            top_p = self._request_metadata["top_p"].tolist()
+
+            for i, (t, k, p) in enumerate(zip(temp, top_k, top_p)):
+                h = (t, k, p)
+                bucket = bucket_map.get(h, None)
+                if bucket is None:
+                    bucket_map[h] = ([i], i)
+                else:
+                    bucket[0].append(i)
 
             # Store the buckets and their equivalence class representatives.
             self._torch_sampling_buckets = (
-                (sampling_buckets[idx], temp_reps[idx], top_k_reps[idx], top_p_reps[idx])
-                for idx in range(len(sampling_buckets))
+                (indices, temp[rep], top_k[rep], top_p[rep])
+                for indices, rep in bucket_map.values()
             )
 
     def _dynamic_step_sample_logits(self):
@@ -652,7 +650,7 @@ def _dynamic_step_sample_logits(self):
                         self._sampling_logits_cuda[indices, :], temp, top_k, top_p
                     )
                 )
-                indices_list.append(indices)
+                indices_list.append(torch.tensor(indices))
 
             # Single write to the output tensor.
             sampled_tokens = torch.cat(token_list, dim=0)