Avoid casting/copying request metadata

Author: tdene

megatron/core/inference/contexts/dynamic_context.py

@@ -538,12 +538,8 @@ def allocate_all_tensors(self, *, is_init: bool) -> None:
 
         # Track request metadata.
         self.request_metadata = {
-            label: torch.empty(
-                (self.max_total_requests,),
-                dtype=dtype,
-                device=torch.cuda.current_device() if on_gpu else torch.device("cpu"),
-            )
-            for label, dtype, on_gpu in self.request_metadata_types
+            label: torch.empty((self.max_total_requests,), dtype=dtype, device=torch.cuda.current_device())
+            for label, dtype, _ in self.request_metadata_types
         }
 
         # Per-token state.
@@ -1054,11 +1050,11 @@ def add_dummy_requests_parallel(
         self.request_output_lengths[request_slice] = lengths_tensor + tokens_to_generate_tensor
         self.request_kv_length_offsets[request_slice] = 0
         self.request_kv_block_counts[request_slice] = block_counts
-        for i, (label, dtype, on_gpu) in enumerate(self.request_metadata_types):
+        for i, (label, dtype, _) in enumerate(self.request_metadata_types):
             self.request_metadata[label][request_slice] = torch.tensor(
                 metadata_cols[i],
                 dtype=dtype,
-                device=torch.cuda.current_device() if on_gpu else torch.device("cpu"),
+                device=torch.cuda.current_device(),
             )
 
         dummy_block_idx = self.block_allocator.dummy_block_idx

megatron/core/inference/inference_request.py

@@ -326,13 +326,13 @@ def get_metadata_types() -> List[Tuple[str, torch.dtype, bool]]:
                 on_device (bool) - Whether the metadata lives on GPU (True) or CPU (False).
         """
         return [
-            ("temperature", torch.float32, True),
-            ("top_k", torch.float32, True),  # CPU for torch sampling
-            ("top_p", torch.float32, True),  # CPU for torch sampling
-            ("termination_id", torch.float32, True),
-            ("return_log_probs", torch.float32, True),  # CPU for non-selective logprobs
-            ("skip_prompt_log_probs", torch.float32, True),  # CPU for non-selective logprobs
-            ("top_n_logprobs", torch.float32, True),  # CPU for torch sampling
+            ("temperature", torch.float32, False),  # CPU for torch sampling
+            ("top_k", torch.int32, False),  # CPU for torch sampling
+            ("top_p", torch.float32, False),  # CPU for torch sampling
+            ("termination_id", torch.int64, True),
+            ("return_log_probs", torch.bool, False),  # CPU for non-selective logprobs
+            ("skip_prompt_log_probs", torch.bool, False),  # CPU for non-selective logprobs
+            ("top_n_logprobs", torch.int32, False), # CPU for torch sampling
         ]
 
     def add_event(self, type: DynamicInferenceEventType, payload: Optional[Any] = None) -> None:

megatron/core/inference/text_generation_controllers/text_generation_controller.py

@@ -90,24 +90,28 @@ def _init_dynamic_sampling_tensors(self):
 
         self._sampling_backend = "torch"
 
-        # Keep track of request metadata.
-        self._active_request_count: int = None
-        self._active_request_slice: slice = None
-        self._request_metadata: Tensor = None
-
         # Initialize bookkeeping tensors.
         self._sampling_logits_cuda = torch.empty(
             max_requests, vocab_size, dtype=logits_dtype, device=device
         )
         self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device)
 
-        self.temperature_cuda = torch.empty_like(self._sampled_tokens_cuda, dtype=torch.float)
-        self.top_k_cuda = torch.empty_like(self._sampled_tokens_cuda, dtype=torch.int32)
-        self.top_p_cuda = torch.empty_like(self._sampled_tokens_cuda, dtype=torch.float)
-        self.termination_id_cuda = torch.empty(max_requests, dtype=torch.int64, device=device)
-        self.return_log_probs_cuda = torch.empty(max_requests, dtype=torch.bool, device=device)
-        self.skip_prompt_log_probs_cuda = torch.empty(max_requests, dtype=torch.bool, device=device)
-        self.top_n_logprobs_cuda = torch.empty(max_requests, dtype=torch.int32, device=device)
+        # Keep track of request metadata.
+        self._request_metadata_raw = {}
+        for label, dtype, on_gpu in context.request_metadata_types:
+            tensor = context.request_metadata[label]
+            if on_gpu:
+                # This is just the same Tensor object as in the context.
+                self._request_metadata_raw[label] = tensor
+            else:
+                # Create pinned tensors for request metadata that lives on CPU.
+                # This is metadata which requires D2H copies, such as top_k for torch sampling.
+                cpu_tensor = torch.empty_like(tensor, device="cpu", pin_memory=True)
+                self._request_metadata_raw[label] = cpu_tensor
+        # These are useful shorthands that are filled in during `_dynamic_step_context_init`.
+        self._active_request_count: int = None
+        self._active_request_slice: slice = None
+        self._request_metadata: Dict[str, Tensor] = {}
 
         # Used for inefficient torch sampling.
         if self._sampling_backend == "torch":
@@ -533,10 +537,21 @@ def _dynamic_step_context_init(
         self._active_request_slice = slice(
             context.paused_request_count, context.total_request_count
         )
-        self._request_metadata = {
-            label: tensor[self._active_request_slice]
-            for label, tensor in context.request_metadata.items()
-        }
+        for label, dtype, on_gpu in context.request_metadata_types:
+            self._request_metadata[label] = self._request_metadata_raw[label][
+                self._active_request_slice
+            ]
+            if on_gpu:
+                # _request_metadata_raw is just a pointer to the Tensor objects in the context.
+                pass
+            if not on_gpu:
+                # We need a D2H copy from the context to the pinned memory buffer.
+                self._request_metadata[label].copy_(
+                    context.request_metadata[label][
+                        self._active_request_slice
+                    ],
+                    non_blocking=True,
+                )
 
         # Get flat tokens, position ids.
         if construct_graph_dimensions is not None:
@@ -596,27 +611,6 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
     def _dynamic_step_sample_bookkeeping(self):
         """Perform bookkeeping necessary to sample logits for dynamic batching."""
 
-        # Copy data into relevant tensors.
-        self.temperature_cuda[:self._active_request_count].copy_(
-            self._request_metadata["temperature"], non_blocking=True
-        )
-        self.top_k_cuda[:self._active_request_count] = self._request_metadata["top_k"].to(
-            dtype=torch.int32, copy=True, non_blocking=True
-        )
-        self.top_p_cuda[:self._active_request_count].copy_(self._request_metadata["top_p"], non_blocking=True)
-        self.termination_id_cuda[:self._active_request_count] = self._request_metadata["termination_id"].to(
-            dtype=torch.int64, copy=True, non_blocking=True
-        )
-        self.return_log_probs_cuda[:self._active_request_count] = self._request_metadata["return_log_probs"].to(
-            dtype=torch.bool, copy=True, non_blocking=True
-        )
-        self.skip_prompt_log_probs_cuda[:self._active_request_count] = self._request_metadata[
-            "skip_prompt_log_probs"
-        ].to(dtype=torch.bool, copy=True, non_blocking=True)
-        self.top_n_logprobs_cuda[:self._active_request_count] = self._request_metadata[
-            "top_n_logprobs"
-        ].to(dtype=torch.int32, copy=True, non_blocking=True)
-
         if self._sampling_backend == "torch":
             # Bucketize the core sampling parameters.
             core_params = torch.stack(
@@ -632,12 +626,10 @@ def _dynamic_step_sample_bookkeeping(self):
             )
             order = torch.argsort(inv_indices, stable=True)
             sampling_buckets = torch.split(order, cnts.tolist())
-            # Perform the D2H sync needed by `_torch_sampling_func` here.
             group_reps = torch.stack([indices[0] for indices in sampling_buckets], dim=0)
-            core_params_reps = core_params[group_reps].detach().cpu()
-            temp_reps = core_params_reps[:, 0].tolist()
-            top_k_reps = core_params_reps[:, 1].to(torch.int32).tolist()
-            top_p_reps = core_params_reps[:, 2].tolist()
+            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()
 
             # Store the buckets and their equivalence class representatives.
             self._torch_sampling_buckets = (
@@ -665,17 +657,17 @@ def _dynamic_step_sample_logits(self):
             # Single write to the output tensor.
             sampled_tokens = torch.cat(token_list, dim=0)
             sampled_indices = torch.cat(indices_list, dim=0)
-            self._sampled_tokens_cuda.index_copy_(0, sampled_indices, sampled_tokens)
+            self._sampled_tokens_cuda[sampled_indices] = sampled_tokens
 
     def _dynamic_step_log_probs_bookkeeping(self) -> Tuple[bool, bool]:
         """Perform bookkeeping necessary to compute log probs for dynamic batching.
 
         Returns:
             return_log_probs (bool): Whether to return the sampled log_probs.
         """
-        return_log_probs = self.return_log_probs_cuda[:self._active_request_count]
-        skip_prompt_log_probs = self.skip_prompt_log_probs_cuda[:self._active_request_count]
-        top_n_log_probs = self.top_n_logprobs_cuda[:self._active_request_count] > 0
+        return_log_probs = self._request_metadata["return_log_probs"]
+        skip_prompt_log_probs = self._request_metadata["skip_prompt_log_probs"]
+        top_n_log_probs = self._request_metadata["top_n_logprobs"] > 0
 
         to_check_prompt = (return_log_probs | top_n_log_probs) & ~skip_prompt_log_probs
 
@@ -727,7 +719,7 @@ def _dynamic_step_calculate_top_n_logprobs(
 
             top_n_results = {}
             for req_idx in range(self._active_request_count):
-                top_n = int(self.top_n_logprobs_cuda[req_idx].item())
+                top_n = int(self._request_metadata["top_n_logprobs"][req_idx].item())
                 if top_n > 0:
                     # Get top-n logprobs and indices for this request (single token)
                     top_n_logits = torch.topk(log_probs[req_idx], k=top_n)
@@ -751,12 +743,12 @@ def _dynamic_step_calculate_top_n_logprobs(
 
         top_n_results = {}
         for req_idx in range(self._active_request_count):
-            top_n = int(self.top_n_logprobs_cuda[req_idx].item())
+            top_n = int(self._request_metadata["top_n_logprobs"][req_idx].item())
             if top_n > 0:
                 request_log_probs = log_probs_per_request[
                     req_idx
                 ]  # [num_tokens_for_request, vocab_size]
-                skip_prompt = bool(self.skip_prompt_log_probs_cuda[req_idx].item())
+                skip_prompt = bool(self._request_metadata["skip_prompt_log_probs"][req_idx].item())
 
                 # If skip_prompt_log_probs is True, only compute for last token
                 if skip_prompt and request_log_probs.size(0) > 1:
@@ -780,6 +772,12 @@ def _dynamic_step_calculate_top_n_logprobs(
     def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]:
         """Update the dynamic inference context after sampling.
 
+        Args:
+            new_sample (Tensor): The newly sampled tokens.
+            request_metadata (Optional[Dict[str, Tensor]]): An override for the tensors
+                that manage request metadata, such as sampling parameters. By default, this
+                metadata is retrieved from the context.
+
         Return:
             Dict [str, Tensor]: A dictionary containing:
                 active_request_ids (Tensor): Current active request IDs.
@@ -799,8 +797,7 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]:
         # Request finished if termination_id or length >= max_sequence_length.
         # Note: termination_id tensor has per-request termination IDs from mixed sampling
         active_request_mask = (
-            self._sampled_tokens_cuda[:self._active_request_count]
-            != self.termination_id_cuda[:self._active_request_count]
+            self._sampled_tokens_cuda[:self._active_request_count] != self._request_metadata["termination_id"]
         ).byte() & torch.less(active_sequence_lengths, max_sequence_lengths).byte()
         finished_idxs = (
             torch.nonzero(active_request_mask == 0, as_tuple=True)[0] + context.paused_request_count