diff --git a/megatron/core/inference/data_parallel_inference_coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator.py index 50f586cc598..aa67e11974d 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator.py +++ b/megatron/core/inference/data_parallel_inference_coordinator.py @@ -586,10 +586,6 @@ def detokenize(self, finished_request): finished_request (dict): The serialized merged request containing the generated tokens to be detokenized. It is modified in place. """ - if finished_request["prompt"] is None: - finished_request["prompt"] = TextGenerationController.detokenize( - self.tokenizer, finished_request["prompt_tokens"][1], remove_EOD=False - ) detokenize_stop_sequence = (finished_request.get("sampling_params", {}) or {}).get( "detokenize_stop_sequence", False ) diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py index 96de2e19713..bb82c73ada8 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py @@ -323,22 +323,6 @@ def _sanitize_tools_for_template(tools): return sanitized -def _reconstruct_reasoning_content(messages: list[dict]) -> list[dict]: - """Reconstruct tags from reasoning_content fields on assistant messages. - - For parity with vLLM, assistant messages may carry reasoning in the reasoning_content field. - Before applying the chat template, we must inline those tags back into content. - """ - for message in messages: - if message.get("role") != "assistant": - continue - reasoning_content = message.pop("reasoning_content", None) - if reasoning_content is not None: - content = message.get("content") or "" - message["content"] = f"{reasoning_content}{content}" - return messages - - def _replace_prefix_tokens( eos_token_id, previous_turn_token_ids, @@ -462,7 +446,6 @@ async def chat_completions(): if not isinstance(messages, list): return Response("'messages' must be a list", status=400) template_messages = _sanitize_messages_for_template(messages) - template_messages = _reconstruct_reasoning_content(template_messages) template_tools = _sanitize_tools_for_template(tools) try: @@ -592,6 +575,7 @@ async def chat_completions(): prompt_tokens = [tokenizer.bos] + prompt_tokens max_tokens = req.get("max_completion_tokens", None) or req.get("max_tokens", None) + ignore_eos = bool(req.get("ignore_eos", False)) sampling_params = SamplingParams( temperature=temperature, @@ -602,6 +586,7 @@ async def chat_completions(): num_tokens_to_generate=(int(max_tokens) if max_tokens is not None else None), skip_prompt_log_probs=skip_prompt_log_probs, add_BOS=add_BOS, + termination_id=-1 if ignore_eos else None, ) except ValueError as e: return Response(f"Invalid sampling parameter: {e}", status=400) @@ -663,6 +648,14 @@ async def chat_completions(): total_completion_tokens = 0 prompt_tokens_counts = [] + prevent_retokenization = req.get("prevent_retokenization", True) + # return_tokenized_data controls whether prompt/generation token ids are + # included in the response. It is independent of prevent_retokenization + # (a client may want token ids without prevent_retokenization, or vice versa), + # but prevent_retokenization implicitly requires token ids so the client + # can echo them back next turn. + return_tokenized_data = req.get("return_tokenized_data", False) or prevent_retokenization + return_raw_text = req.get("return_raw_text", False) request_idx = 0 for result_item in batch_results: result = unwrap_serialized_tensors(result_item) @@ -736,9 +729,14 @@ async def chat_completions(): if "reasoning" in metadata: message["reasoning_content"] = metadata["reasoning"] - # Replicate data in the message field for compatibility. - message["prompt_token_ids"] = result["prompt_tokens"] - message["generation_token_ids"] = result["generated_tokens"] + if return_tokenized_data: + message["prompt_token_ids"] = result["prompt_tokens"] + message["generation_token_ids"] = result["generated_tokens"] + if return_raw_text: + prompt_str = tokenizer.detokenize(result["prompt_tokens"]) + message["raw_text"] = prompt_str + text_output + # Small RL/debug scalars (a few bytes each); harmless to keep for + # NeMo-RL compatibility. message["generation_log_probs"] = result.get("generated_log_probs", []) message["policy_epoch"] = result["policy_epoch"] message["kv_cache_epoch"] = result["kv_cache_epoch"] @@ -759,15 +757,13 @@ async def chat_completions(): else: finish_reason = "stop" + # Choice-level prompt/generation_token_ids, generation_log_probs and + # raw_text were duplicates of message-level data (or reconstructable); + # dropped to match vLLM's response shape and cut payload size. choice_data = { "index": request_idx, "message": message, - "prompt_token_ids": result["prompt_tokens"], - "generation_token_ids": result["generated_tokens"], - "generation_log_probs": result.get("generated_log_probs", []), - "raw_text": result["prompt"] + result["generated_text"], # 'logprobs' in chat API is an object containing 'content' - # "logprobs": {"content": logprobs_content} if logprobs_content else None, "logprobs": {"content": logprobs_content} if return_log_probs else None, "finish_reason": finish_reason, } @@ -782,7 +778,7 @@ async def chat_completions(): ] choices.append(choice_data) - if choice_data["generation_log_probs"] is None: + if result.get("generated_log_probs") is None: logger.warning( "Generation log probs is None for request:\n%s", json.dumps(_redact_token_id_lists_for_logging(result), indent=4), diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py index 055b232d73f..b2a6b320a97 100644 --- a/megatron/rl/inference/megatron.py +++ b/megatron/rl/inference/megatron.py @@ -65,6 +65,12 @@ async def base_generate(self, request: InferenceRequest) -> InferenceResponse: extra_body={ "skip_prompt_log_probs": True, "add_BOS": (not args.rl_skip_bos_token and tokenizer.bos is not None), + # TODO: These are non-standard fields that add significant memory overheads to the + # chat completions payload. return_raw_text also wastes a lot of CPU cycles + # detokenizing prompt tokens, especially expensive for long prompts in agentic RL. + # Set to False if not needed in MRL. + "return_tokenized_data": True, + "return_raw_text": True, }, ) @@ -73,11 +79,11 @@ async def base_generate(self, request: InferenceRequest) -> InferenceResponse: return InferenceResponse( # TODO: Handle tool calls and reasoning in LLMChatMessage response=LLMChatMessage(**choice.message.model_dump(include={'role', 'content'})), - raw_text=choice.raw_text, - token_ids=choice.prompt_token_ids + choice.generation_token_ids, - logprobs=choice.generation_log_probs, + raw_text=choice.message.raw_text, + token_ids=choice.message.prompt_token_ids + choice.message.generation_token_ids, + logprobs=choice.message.generation_log_probs, finish_reason=choice.finish_reason, - prompt_length=len(choice.prompt_token_ids), + prompt_length=len(choice.message.prompt_token_ids), policy_epoch=choice.message.policy_epoch, kv_cache_epoch=choice.message.kv_cache_epoch, num_evictions=choice.message.num_evictions,