Optimizing post-processing of requests (#2920)

Author: sidsingh-nvidia

megatron/core/inference/data_parallel_inference_coordinator.py

@@ -65,7 +65,7 @@ class DataParallelInferenceCoordinator:
         next_request_id (int): A counter for generating unique server-side request IDs.
     """
 
-    def __init__(self, inference_coordinator_port: int, data_parallel_size: int):
+    def __init__(self, inference_coordinator_port: int, data_parallel_size: int, tokenizer):
         """
         Initializes the inference coordinator.
 
@@ -116,6 +116,7 @@ def __init__(self, inference_coordinator_port: int, data_parallel_size: int):
         self.request_id_to_client_request_id = {}
 
         self.next_request_id = 0
+        self.tokenizer = tokenizer
 
     def get_next_data_parallel_rank(self):
         """
@@ -261,6 +262,7 @@ def start(self):
                 finished_request_records = deserialized_payload[1]
 
                 for finished_request_record in finished_request_records:
+                    self.detokenize(finished_request_record)
                     fid = finished_request_record["requests"][0]["request_id"]
                     client_identity = self.request_id_to_client_id[fid]
                     client_request_identity = self.request_id_to_client_request_id[fid]
@@ -280,9 +282,25 @@ def start(self):
             else:
                 raise UnknownHeaderError(header)
 
+    def detokenize(self, finished_request_record):
+        """
+        Detokenizes the generated tokens in the finished request record.
+
+        This method uses the coordinator's tokenizer to convert the list of
+        generated token IDs back into human-readable text.
+
+        Args:
+            finished_request_record (dict): The record containing the generated
+                tokens to be detokenized. It is modified in place.
+        """
+        for request in finished_request_record["requests"]:
+            if request["prompt"] is None:
+                request["prompt"] = self.tokenizer.detokenize(request["prompt_tokens"][1])
+            request["generated_text"] = self.tokenizer.detokenize(request["generated_tokens"])
+
     @classmethod
     def entrypoint(
-        cls, ready_event: Event, inference_coordinator_port: int, data_parallel_size: int
+        cls, ready_event: Event, inference_coordinator_port: int, data_parallel_size: int, tokenizer
     ):
         """
         Class method to instantiate and run the coordinator, for use in a separate process.
@@ -296,7 +314,7 @@ def entrypoint(
             inference_coordinator_port (int): The port to bind to.
             data_parallel_size (int): The number of expected TP-coordinators.
         """
-        coordinator = cls(inference_coordinator_port, data_parallel_size)
+        coordinator = cls(inference_coordinator_port, data_parallel_size, tokenizer)
         ready_event.set()
         try:
             coordinator.start()

megatron/core/inference/engines/dynamic_engine.py

@@ -413,6 +413,7 @@ async def start_listening_to_data_parallel_coordinator(
                     coordinator_ready_event,
                     inference_coordinator_port,
                     get_pg_size(self.pg_collection.dp),
+                    self.controller.tokenizer,
                 ),
             )
             self.inference_coordinator_process.start()
@@ -1205,6 +1206,7 @@ async def async_bookkeep(
                 cuda_graph_request_count (int): The CUDA graph batch size matching this step.
         """
         # Increment finished_request_count.
+        range_push("bookkeeping")
         cuda_graph_request_count = None
 
         if step_result is not None:
@@ -1248,26 +1250,33 @@ async def async_bookkeep(
             finished_request_records.append(failed_entry.record)
             failed_entry.future.set_result(failed_entry.record)
         self.failed_request_ids.clear()
+        range_pop()
 
-        # Detokenize all finished requests (critical for InferenceClient, which
-        # doesn't necessarily have the tokenizer).
-        for record in finished_request_records:
-            for request in record.requests:
-                if request.prompt is None:
-                    request.prompt = self.controller.tokenizer.detokenize(
-                        request.prompt_tokens.tolist()
+        # Detokenize all finished requests if not using
+        # the coordinator. Otherwise, the coordinator will
+        # overlap detokenization with the engine.
+        if not self.use_coordinator:
+            range_push("detokenization")
+            for record in finished_request_records:
+                for request in record.requests:
+                    if request.prompt is None:
+                        request.prompt = self.controller.tokenizer.detokenize(
+                            request.prompt_tokens.tolist()
+                        )
+                    request.generated_text = self.controller.tokenizer.detokenize(
+                        request.generated_tokens
                     )
-                request.generated_text = self.controller.tokenizer.detokenize(
-                    request.generated_tokens
-                )
+            range_pop()
 
         # Handle necessary ZMQ DP coordinator communication.
         if self.use_coordinator and self.is_mp_coordinator and finished_request_records:
+            range_push("coordinator_communication")
             payload = msgpack.packb(
                 [Headers.ENGINE_REPLY.value, [r.serialize() for r in finished_request_records]],
                 use_bin_type=True,
             )
             self.socket_for_receiving_requests.send(payload)
+            range_pop()
 
         # Log KV cache utilization stats to W&B
         if context_state["kv_stats"] is not None:
@@ -1461,7 +1470,7 @@ def schedule_requests(self) -> int:
             int: The number of messages that were received and processed in this batch.
         """
 
-        torch.cuda.nvtx.range_push("drain_zmq_socket")
+        range_push("drain_zmq_socket")
         all_messages = []
         if self.is_mp_coordinator:
             while True:
@@ -1494,7 +1503,7 @@ def schedule_requests(self) -> int:
             else:
                 all_messages = []
 
-        torch.cuda.nvtx.range_pop()
+        range_pop()
         for message in all_messages:
             data = msgpack.unpackb(message, raw=False)
             header = Headers(data[0])
@@ -1507,7 +1516,9 @@ def schedule_requests(self) -> int:
             if header == Headers.SUBMIT_REQUEST:
                 request_id, prompt, sampling_params = data[1:]
                 sampling_params = SamplingParams.deserialize(sampling_params)
+                range_push("add_request")
                 self.add_request(request_id, prompt, sampling_params)
+                range_pop()
             elif header == Headers.PAUSE:
                 # Pause thyself.
                 self.received_pause = True

megatron/core/inference/inference_request.py

@@ -1,7 +1,6 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import copy
-import io
 import time
 import warnings
 from dataclasses import asdict, dataclass, field
@@ -15,33 +14,34 @@
 from megatron.core.utils import experimental_api
 
 
-def serialize_tensor(tensor: torch.Tensor) -> bytes:
+def serialize_tensor(tensor: torch.Tensor) -> List:
     """Serialize tensor to bytes.
 
     Args:
         tensor (Tensor): Tensor.
 
     Returns:
-        (bytes) Byte representation of tensor.
+        (List) Tensor as a list
     """
-    buffer = io.BytesIO()
-    torch.save(tensor, buffer)
-    buffer.seek(0)
-    tensor_bytes = buffer.read()
-    return tensor_bytes
+    torch.cuda.nvtx.range_push("serialize_tensor")
 
+    # simply convert tensor into a list
+    tensor = tensor.cpu().tolist()
 
-def deserialize_tensor(tensor_bytes: bytes) -> torch.Tensor:
+    torch.cuda.nvtx.range_pop()
+    return tensor
+
+
+def deserialize_tensor(tensor_as_list: List) -> torch.Tensor:
     """Deserialize tensor from bytes.
 
     Args:
-        tensor_bytes (bytes): Byte representation of tensor.
+        tensor_as_list (List): List representation of tensor.
 
     Returns:
         (Tensor) Tensor.
     """
-    buffer = io.BytesIO(tensor_bytes)
-    tensor = torch.load(buffer)
+    tensor = torch.tensor(tensor_as_list)
     return tensor
 
 
@@ -99,17 +99,21 @@ def serialize(self) -> dict:
             (dict) A dictionary representation of the instance suitable for
                 serialization.
         """
-
         # Dataclass to dict.
-        obj = asdict(self)
+        # do not use asdict(self) - it has very high CPU overheads
+        # and if there are tensors, it will try to deepcopy them
+        obj = self.__dict__.copy()  # shallow dict copy
         obj["status"] = self.status.name if self.status else None
+        obj["sampling_params"] = self.sampling_params.serialize() if self.sampling_params else None
+        obj["inference_parameters"] = (
+            self.inference_parameters.serialize() if self.inference_parameters else None
+        )
 
         # Serialize tensors.
         obj = {
             k: (("tensor", serialize_tensor(v)) if isinstance(v, torch.Tensor) else v)
             for k, v in obj.items()
         }
-
         return obj
 
     @classmethod
@@ -125,14 +129,31 @@ def deserialize(cls, obj: dict) -> "InferenceRequest":
 
         # Initialize request.
         request = cls(**obj)
-        request.status = None if obj["status"] is None else Status[obj["status"]]
+        request._post_deserialize(obj)
+        return request
 
-        # Deserialize tensors.
+    def _post_deserialize(self, obj: dict):
+        """
+        This is called after the dataclass is initialized to handle any special
+        deserialization logic.
+        """
+        # Deserialize status.
+        self.status = None if obj["status"] is None else Status[obj["status"]]
+        self.sampling_params = (
+            None
+            if obj["sampling_params"] is None
+            else SamplingParams.deserialize(obj["sampling_params"])
+        )
+        self.inference_parameters = (
+            None
+            if obj["inference_parameters"] is None
+            else SamplingParams.deserialize(obj["inference_parameters"])
+        )
+
+        # Deserialize tensors and sampling params.
         for k, v in obj.items():
             if isinstance(v, list) and len(v) == 2 and v[0] == "tensor":
-                setattr(request, k, deserialize_tensor(v[1]))
-
-        return request
+                setattr(self, k, deserialize_tensor(v[1]))
 
 
 class DynamicInferenceEventType(Enum):
@@ -197,15 +218,18 @@ def serialize(self) -> dict:
         """
 
         # Dataclass to dict.
-        obj = asdict(self)
+        torch.cuda.nvtx.range_push("DynamicInferenceEvent.serialize")
+        # do not use asdict(self) - it has very high CPU overheads
+        # and if there are tensors, it will try to deepcopy them
+        obj = self.__dict__.copy()
         obj["type"] = self.type.name
 
         # Serialize payload.
         if self.payload:
             from .contexts.dynamic_context import ContextErrorFactory  # avoid circular import.
 
             obj["payload"] = ContextErrorFactory.serialize(self.payload)
-
+        torch.cuda.nvtx.range_pop()
         return obj
 
     @classmethod
@@ -247,7 +271,7 @@ class DynamicInferenceRequest(InferenceRequest):
     # remaining prompt tokens are used for chunked prefill
     remaining_prompt_tokens: Optional[torch.Tensor] = None
     latency: Optional[float] = None
-    finished_chunk_token_count = 0
+    finished_chunk_token_count: int = 0
     stop_word_ids: Optional[List[List[int]]] = None  # Tokenized stop words (populated internally)
 
     def __post_init__(self):
@@ -275,30 +299,22 @@ def __str__(self):
             )
         )
 
-    def serialize(self):
+    def serialize(self) -> dict:
         """Converts the instance into a serializable dictionary.
 
         Returns:
             (dict) A dictionary representation of the instance suitable for
                 serialization.
         """
+        torch.cuda.nvtx.range_push("DynamicInferenceRequest.serialize")
         obj = super().serialize()
         obj["events"] = [e.serialize() for e in self.events]
+        torch.cuda.nvtx.range_pop()
         return obj
 
-    @classmethod
-    def deserialize(cls, obj: dict) -> "DynamicInferenceRequest":
-        """Deserialize request.
-
-        Args:
-            obj (dict): Serialized request data.
-
-        Returns:
-            (DynamicInferenceRequest) Deserialized request.
-        """
-        request = super().deserialize(obj)
-        request.events = [DynamicInferenceEvent.deserialize(e) for e in obj["events"]]
-        return request
+    def _post_deserialize(self, obj):
+        super()._post_deserialize(obj)
+        self.events = [DynamicInferenceEvent.deserialize(e) for e in obj["events"]]
 
     @property
     def tracked_metadata(self) -> List[Any]:
@@ -517,8 +533,10 @@ def serialize(self) -> dict:
             (dict) A dictionary representation of the instance suitable for
                 serialization.
         """
-        obj = asdict(self)
-        obj["requests"] = [r.serialize() for r in self.requests]
+        torch.cuda.nvtx.range_push("DynamicInferenceRequestRecord.serialize")
+        obj = self.__dict__.copy()  # shallow dict copy
+        obj["requests"] = [r.serialize() for r in obj["requests"]]
+        torch.cuda.nvtx.range_pop()
         return obj
 
     @classmethod