lint

Author: tdene

megatron/core/inference/contexts/dynamic_context.py

@@ -538,7 +538,9 @@ 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())
+            label: torch.empty(
+                (self.max_total_requests,), dtype=dtype, device=torch.cuda.current_device()
+            )
             for label, dtype, _ in self.request_metadata_types
         }
 
@@ -1052,9 +1054,7 @@ def add_dummy_requests_parallel(
         self.request_kv_block_counts[request_slice] = block_counts
         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(),
+                metadata_cols[i], dtype=dtype, device=torch.cuda.current_device()
             )
 
         dummy_block_idx = self.block_allocator.dummy_block_idx

megatron/core/inference/inference_request.py

@@ -12,6 +12,7 @@
 
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.tokenizers import MegatronTokenizer
+from megatron.core.utils import internal_api
 
 
 def serialize_tensor(tensor: torch.Tensor) -> bytes:
@@ -228,6 +229,7 @@ def deserialize(cls, obj: dict) -> "DynamicInferenceEvent":
         return event
 
 
+@internal_api
 @dataclass(kw_only=True)
 class DynamicInferenceRequest(InferenceRequest):
     """Class for one inference request
@@ -332,7 +334,7 @@ def get_metadata_types() -> List[Tuple[str, torch.dtype, bool]]:
             ("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
+            ("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

@@ -20,7 +20,7 @@
     is_pipeline_last_stage,
 )
 from megatron.core.inference.contexts.dynamic_context import MaxSequenceLengthOverflowError
-from megatron.core.inference.inference_request import DynamicInferenceRequest, InferenceRequest, Status
+from megatron.core.inference.inference_request import InferenceRequest, Status
 from megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper import (
     AbstractModelInferenceWrapper,
 )
@@ -547,10 +547,7 @@ def _dynamic_step_context_init(
             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,
+                    context.request_metadata[label][self._active_request_slice], non_blocking=True
                 )
 
         # Get flat tokens, position ids.
@@ -604,7 +601,9 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
         else:
             last_token_logits = context.last_token_logits(logits)
         # Copy last_token_logits to contiguous buffer.
-        self._sampling_logits_cuda[:self._active_request_count].copy_(last_token_logits, non_blocking=True)
+        self._sampling_logits_cuda[: self._active_request_count].copy_(
+            last_token_logits, non_blocking=True
+        )
 
         return logits
 
@@ -631,8 +630,7 @@ def _dynamic_step_sample_bookkeeping(self):
 
             # Store the buckets and their equivalence class representatives.
             self._torch_sampling_buckets = (
-                (indices, temp[rep], top_k[rep], top_p[rep])
-                for indices, rep in bucket_map.values()
+                (indices, temp[rep], top_k[rep], top_p[rep]) for indices, rep in bucket_map.values()
             )
 
     def _dynamic_step_sample_logits(self):
@@ -683,7 +681,7 @@ def _dynamic_step_calculate_log_probs(self, logits: Tensor) -> Optional[Tensor]:
 
         return context.calculate_log_probs(
             logits,
-            self._sampled_tokens_cuda[:self._active_request_count],
+            self._sampled_tokens_cuda[: self._active_request_count],
             only_last_token_logits=self._materialize_only_last,
         )
 
@@ -713,7 +711,7 @@ def _dynamic_step_calculate_top_n_logprobs(
         if self._materialize_only_last or context.is_decode_only():
             # In decode mode or when only last token logits are materialized,
             # logits already represent only the last tokens
-            log_probs = log_probs_tensor[:self._active_request_count]
+            log_probs = log_probs_tensor[: self._active_request_count]
 
             top_n_results = {}
             for req_idx in range(self._active_request_count):
@@ -731,9 +729,7 @@ def _dynamic_step_calculate_top_n_logprobs(
         # Note: logits may be padded, so we only take the first active_token_count tokens
         log_probs = log_probs_tensor[: context.active_token_count]
 
-        active_query_lengths = context.request_query_lengths[
-            self._active_request_slice
-        ]
+        active_query_lengths = context.request_query_lengths[self._active_request_slice]
 
         # Split log_probs across request boundaries
         # log_probs has shape [active_token_count, vocab_size]
@@ -785,25 +781,24 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]:
         context = self.inference_wrapped_model.inference_context
 
         # Active sequence lengths.
-        active_request_ids = context.request_ids[
-            self._active_request_slice
-        ].long()
+        active_request_ids = context.request_ids[self._active_request_slice].long()
         active_sequence_lengths = context.get_active_sequence_lengths()
         active_sequence_lengths += 1  # Account for the token we just generated
         max_sequence_lengths = context.get_max_sequence_lengths()
 
         # 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._request_metadata["termination_id"]
+            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
         )
         finished_request_ids = context.request_ids[finished_idxs]
 
         # New sample gets updated in update_requests, so we pass in a clone
-        new_sample_copy = self._sampled_tokens_cuda[:self._active_request_count].clone()
+        new_sample_copy = self._sampled_tokens_cuda[: self._active_request_count].clone()
 
         # Update requests.
         newly_paused_request_ids = context.update_requests(active_request_mask, new_sample_copy)
@@ -875,7 +870,7 @@ async def async_generate_output_tokens_dynamic_batch(
             request_bookkeeping = self._dynamic_step_context_bookkeeping()
 
         ret = {
-            "sample": self._sampled_tokens_cuda[:self._active_request_count],
+            "sample": self._sampled_tokens_cuda[: self._active_request_count],
             "log_probs": log_probs,
             "top_n_logprobs": top_n_logprobs,
             "cuda_graph_request_count": cuda_graph_request_count,