diff --git a/megatron/core/inference/data_parallel_inference_coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator.py index 7f81065bcd4..3a1747facb2 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator.py +++ b/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() diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index c56f91bbbe9..c42246d5624 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/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 diff --git a/megatron/core/inference/inference_request.py b/megatron/core/inference/inference_request.py index 8bd0dd0aff4..6a7354220f9 100644 --- a/megatron/core/inference/inference_request.py +++ b/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,7 +218,10 @@ 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. @@ -205,7 +229,7 @@ def serialize(self) -> dict: 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