diff --git a/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py index 67eb1bb5829..ac590bdb8ca 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py +++ b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py @@ -464,6 +464,9 @@ 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 not finished_request["sampling_params"]["detokenize_generations"]: + return + detokenize_stop_sequence = (finished_request.get("sampling_params", {}) or {}).get( "detokenize_stop_sequence", False ) diff --git a/megatron/core/inference/sampling_params.py b/megatron/core/inference/sampling_params.py index 1e3a43f0e89..cfdeb477256 100644 --- a/megatron/core/inference/sampling_params.py +++ b/megatron/core/inference/sampling_params.py @@ -38,6 +38,11 @@ class SamplingParams: # drops prompt_tokens before serializing the finished request, saving the ZMQ # transmission cost for long prompts. Opt in when the client needs them. return_prompt_tokens: bool = False + # When False, the DP coordinator skips detokenizing this request's output. The + # coordinator is a single process serving every DP rank, so per-request work there + # is a throughput ceiling; callers that can detokenize themselves (e.g. the HTTP + # frontend, which is replicated) should set this to False. + detokenize_generations: bool = True streaming: bool = False # Emit incremental ENGINE_REPLY_PARTIAL frames. streaming_interval: int = 1 # Minimum unsent tokens per ENGINE_REPLY_PARTIAL. do_kv_handoff: bool = False # Pin KV blocks and expose metadata for peer transfer. 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 ce1550aa55d..6c2c96f3a59 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 @@ -7,9 +7,13 @@ import traceback import uuid import warnings +from functools import partial from megatron.core.inference.inference_request import unwrap_serialized_tensors from megatron.core.inference.sampling_params import SamplingParams +from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + TextGenerationController, +) from megatron.core.tokenizers.text.parsers import PARSER_MAPPING from ..incremental_detokenizer import HuggingFaceFastIncrementalDetokenizer @@ -355,6 +359,25 @@ def _replace_prefix_tokens( return previous_turn_token_ids + current_turn_additional_token_ids +def _apply_chat_template_sync( + tokenizer, messages, tools, chat_template_kwargs, add_generation_prompt=True +): + """Apply the chat template and coerce to `list[int]`, for use in a worker thread. + + The coercion runs here too: it walks every token, so leaving it on the event loop + would keep part of the stall this offload exists to remove. + """ + return _coerce_to_token_id_list( + tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=add_generation_prompt, + tools=tools, + **chat_template_kwargs, + ) + ) + + def _coerce_to_token_id_list(result): """Convert the return value of `tokenizer.apply_chat_template` to `list[int]`. @@ -476,14 +499,15 @@ async def chat_completions(): hasattr(tokenizer, 'apply_chat_template') and getattr(tokenizer, "chat_template", None) is not None ): - prompt_tokens = _coerce_to_token_id_list( - tokenizer.apply_chat_template( + prompt_tokens = await asyncio.get_running_loop().run_in_executor( + current_app.config['tokenize_executor'], + partial( + _apply_chat_template_sync, + current_app.config['tokenize_tokenizer'], template_messages, - tokenize=True, - add_generation_prompt=True, - tools=template_tools, - **chat_template_kwargs, - ) + template_tools, + chat_template_kwargs, + ), ) if req.get("prevent_retokenization", True): @@ -524,13 +548,17 @@ async def chat_completions(): ] # Get the templated tokenization of just the previous generation - retokenized_previous_turn_token_ids = _coerce_to_token_id_list( - tokenizer.apply_chat_template( - messages_to_last_assistant_message, - tokenize=True, - add_generation_prompt=False, - tools=template_tools, - **chat_template_kwargs, + retokenized_previous_turn_token_ids = ( + await asyncio.get_running_loop().run_in_executor( + current_app.config['tokenize_executor'], + partial( + _apply_chat_template_sync, + current_app.config['tokenize_tokenizer'], + messages_to_last_assistant_message, + template_tools, + chat_template_kwargs, + add_generation_prompt=False, + ), ) ) @@ -623,6 +651,9 @@ async def chat_completions(): termination_id=-1 if ignore_eos else None, return_prompt_tokens=return_prompt_tokens, streaming_interval=int(_get_non_none(req, "streaming_interval", 1)), + # This frontend detokenizes its own output below. Keeping it off the + # coordinator matters because that is one process shared by all DP ranks. + detokenize_generations=False, ) except ValueError as e: return Response(f"Invalid sampling parameter: {e}", status=400) @@ -759,7 +790,11 @@ def parse_streaming_text(text): if response_uid is None: response_uid = result["uid"] - text_output = result["generated_text"] + text_output = TextGenerationController.detokenize( + tokenizer, + result["generated_tokens"], + remove_EOD=not sampling_params.detokenize_stop_sequence, + ) # The engine always reports prompt_length (for usage), but drops the # prompt_tokens tensor unless return_prompt_tokens was set. prompt_tokens_count = result.get("prompt_length") diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py index af12bc5bbc1..2ba3cdbbb7e 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py @@ -1,9 +1,11 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. import asyncio +import copy import logging import multiprocessing as mp import socket +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from typing import List, Optional @@ -24,7 +26,6 @@ # Global reference to manage the background server processes _SERVER_PROCESSES: List[mp.Process] = [] -_SHARED_SOCKET = None @contextmanager @@ -47,7 +48,6 @@ async def _run_text_gen_server( server_port: int, parsers: Optional[List[str]] = None, verbose: bool = False, - fd: Optional[int] = None, hostname: Optional[str] = None, ): """ @@ -63,6 +63,10 @@ async def _run_text_gen_server( logger.info(f"Rank {rank}: InferenceClient connected.") try: + # Bind what the caller asked for -- None means every interface, which is + # not the single address gethostname() resolves to. The resolved name is + # for the log line only. + bind_host = hostname if hostname is None: try: hostname = socket.gethostname() @@ -81,6 +85,15 @@ async def _run_text_gen_server( app.config['parsers'] = parsers app.config['verbose'] = verbose + # Applying the chat template is synchronous and O(prompt); on the event loop it + # stalls every other request this replica owns, including delivery of responses + # that already finished. One worker is enough - the point is the yield, not + # throughput. The copy is required: HF tokenizers are not thread-safe. + app.config['tokenize_executor'] = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="tokenize" + ) + app.config['tokenize_tokenizer'] = copy.deepcopy(tokenizer) + # Register all blueprints from the 'endpoints' package for endpoint in endpoints.__all__: app.register_blueprint(endpoint) @@ -92,18 +105,20 @@ async def _run_text_gen_server( 2**14 ) # Allow many concurrent streams for HTTP/2 clients. - if fd is not None: - config.bind = [f"fd://{fd}"] - else: - config.bind = [f"{hostname}:{server_port}"] + # Held for this worker's lifetime; closing it would drop the listener. + own_socket = _bind_reuseport_socket(server_port, bind_host) + config.bind = [f"fd://{own_socket.fileno()}"] with temp_log_level(logging.INFO, logger): logger.info(f"Starting text generation server on http://{hostname}:{server_port}") logger.info(f"Using tokenizer: {type(tokenizer)}") logger.info(f"Using parsers: {parsers}") - # Quart is natively ASGI, so we can serve the app directly - await serve(app, config) + try: + # Quart is natively ASGI, so we can serve the app directly + await serve(app, config) + finally: + own_socket.close() finally: # Gracefully shut down the client when the server stops @@ -118,7 +133,6 @@ def _server_process_worker( server_port: int, parsers: Optional[List[str]] = None, verbose: bool = False, - fd: Optional[int] = None, hostname: Optional[str] = None, ): """Synchronous worker function that sets up a new event loop for the separate process.""" @@ -127,7 +141,7 @@ def _server_process_worker( try: loop.run_until_complete( _run_text_gen_server( - coordinator_addr, tokenizer, rank, server_port, parsers, verbose, fd, hostname + coordinator_addr, tokenizer, rank, server_port, parsers, verbose, hostname ) ) except KeyboardInterrupt: @@ -141,6 +155,38 @@ def _server_process_worker( loop.close() +def _bind_reuseport_socket(server_port: int, hostname: Optional[str]) -> socket.socket: + """Bind this worker's own socket on the shared port, with SO_REUSEPORT. + + Unlike inheriting one fd, this gives every worker its own accept queue and + lets the kernel hash each connection's 4-tuple across them. It is a large + improvement on sharing (measured 604x -> 2x spread at 32 replicas) but is + still hashing, not balancing: it cannot see that a replica is already busy, + so the spread is statistical rather than exact. + """ + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + # Required on every socket sharing the port; without it the second bind fails. + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) + sock.bind((hostname if hostname is not None else "0.0.0.0", server_port)) + sock.setblocking(False) + return sock + + +def _reserve_port(hostname: Optional[str]) -> int: + """Pick a free port for the replicas to bind individually. + + Replicas each bind the port themselves, so the parent cannot hold the socket + and hand out its fd; it binds only long enough to learn a free port. The gap + before the replicas bind is a small race with unrelated processes, which is + why an explicit port is preferred when one is available. + """ + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + probe.bind((hostname if hostname is not None else "0.0.0.0", 0)) + return probe.getsockname()[1] + + def start_text_gen_server( coordinator_addr: str, tokenizer, @@ -151,77 +197,88 @@ def start_text_gen_server( num_replicas: int = 4, hostname: Optional[str] = None, sock: Optional[socket.socket] = None, -): - """Start the text generation server.""" +) -> Optional[str]: + """Start the text generation server. + + Every replica binds its own socket on ``server_port`` with SO_REUSEPORT, so + each gets its own accept queue and the kernel spreads connections across + them. Sharing one inherited socket does not balance -- replicas race to + accept from a single queue and whichever is already running keeps winning, + which concentrates most traffic on a handful of them as replica count grows. + + Call this on every rank that should host a frontend. Frontend work (chat + template, detokenize, parsers, JSON) is CPU-bound, so hosting on a single + rank confines it to that rank's CPU allocation and leaves the rest of the + job's cores unused. Each caller gets its own URL back; collecting them and + spreading requests over the result is the caller's business. + + Args: + server_port: Port to listen on. Overridden by ``sock`` when given; 0 + asks the OS to choose a free one. + sock: A socket the caller already bound, used only to fix the port. + Replicas bind that port themselves, so it is closed here rather than + shared with them. + + Returns: + The base URL this rank serves on, or None if the server was already + running. + """ global _SERVER_PROCESSES - global _SHARED_SOCKET if _SERVER_PROCESSES: logger.warning("Text gen server processes are already running.") - return + return None - # The caller may pass in a socket it has already bound ahead of time. if sock is not None: - bound_port = sock.getsockname()[1] - if bound_port == 0: + # Take the port and release the socket: replicas each bind their own with + # SO_REUSEPORT, which one shared socket cannot provide. + server_port = sock.getsockname()[1] + if server_port == 0: raise ValueError( "socket must be bound to a real port before being passed to start_text_gen_server" ) - _SHARED_SOCKET = sock - server_port = bound_port - _SHARED_SOCKET.setblocking(False) - else: - _SHARED_SOCKET = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - _SHARED_SOCKET.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - - if hasattr(socket, 'SO_REUSEPORT'): - try: - _SHARED_SOCKET.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) - except OSError: - pass - - bind_address = hostname if hostname is not None else "0.0.0.0" - _SHARED_SOCKET.bind((bind_address, server_port)) - _SHARED_SOCKET.setblocking(False) - - _SHARED_SOCKET.set_inheritable(True) - fd = _SHARED_SOCKET.fileno() + sock.close() + elif server_port == 0: + server_port = _reserve_port(hostname) for i in range(num_replicas): p = mp.Process( target=_server_process_worker, - args=(coordinator_addr, tokenizer, rank, server_port, parsers, verbose, fd, hostname), + args=(coordinator_addr, tokenizer, rank, server_port, parsers, verbose, hostname), daemon=True, ) p.start() _SERVER_PROCESSES.append(p) - logger.info(f"Started text gen frontend replica {i+1}/{num_replicas} (PID: {p.pid})") + logger.info( + f"Started text gen frontend replica {i+1}/{num_replicas} " + f"on port {server_port} (PID: {p.pid})" + ) + return f"http://{hostname or socket.gethostname()}:{server_port}" -def stop_text_gen_server(): - """Stop the text generation server.""" - global _SERVER_PROCESSES - global _SHARED_SOCKET - if not _SERVER_PROCESSES: +def _terminate(processes: List[mp.Process], what: str): + """Terminate a group of worker processes, escalating to kill if needed.""" + if not processes: return - - logger.info(f"Terminating {len(_SERVER_PROCESSES)} Text Gen frontend processes...") - - for p in _SERVER_PROCESSES: + logger.info(f"Terminating {len(processes)} {what} processes...") + for p in processes: if p.is_alive(): p.terminate() - - for p in _SERVER_PROCESSES: + for p in processes: p.join(timeout=3) if p.is_alive(): p.kill() p.join() - # Clean up the master socket - if _SHARED_SOCKET is not None: - _SHARED_SOCKET.close() - _SHARED_SOCKET = None +def stop_text_gen_server(): + """Stop this rank's frontend replica processes.""" + global _SERVER_PROCESSES + + if not _SERVER_PROCESSES: + return + + _terminate(_SERVER_PROCESSES, "Text Gen frontend") _SERVER_PROCESSES = [] logger.info("All text gen frontend processes terminated.") diff --git a/tools/run_dynamic_text_generation_server.py b/tools/run_dynamic_text_generation_server.py index a33b066e8ea..8aca26ba537 100644 --- a/tools/run_dynamic_text_generation_server.py +++ b/tools/run_dynamic_text_generation_server.py @@ -2,6 +2,7 @@ import argparse import asyncio +import logging import torch @@ -10,7 +11,7 @@ start_text_gen_server, stop_text_gen_server, ) -from megatron.core.utils import configure_nvtx_profiling, trace_async_exceptions +from megatron.core.utils import configure_nvtx_profiling, get_pg_size, trace_async_exceptions from megatron.inference.utils import add_inference_args, get_dynamic_inference_engine from megatron.post_training.arguments import add_modelopt_args from megatron.training import get_args @@ -30,6 +31,23 @@ def add_text_generation_server_args(parser: argparse.ArgumentParser): parser.add_argument( "--parsers", type=str, nargs="+", default=[], help="Parsers to use for parsing the response" ) + parser.add_argument( + "--frontend-replicas", type=int, default=-1, + help="Number of HTTP frontend processes spawned per hosting rank. " + "-1 (default) uses max(data parallel size, 4), or a flat 4 with " + "--frontend-on-all-ranks, where capacity already scales with the " + "number of ranks hosting a frontend.", + ) + parser.add_argument( + "--frontend-on-all-ranks", action="store_true", + help="Run HTTP frontends on every rank instead of only rank 0, and " + "return every rank's URL for the caller to spread requests over. " + "Frontend work (chat template, detokenize, parsers, JSON) is " + "CPU-bound and otherwise confined to the hosting rank's CPU " + "allocation, which leaves the rest of the job's cores unused. " + "Ranks still share one DP coordinator; only the HTTP tier is " + "replicated.", + ) return parser @@ -54,25 +72,62 @@ async def run_text_generation_server( hostname=hostname, ) + num_replicas = args.frontend_replicas + if num_replicas < 0: + if args.frontend_on_all_ranks: + # Capacity now scales with the number of ranks, so the per-rank + # replica count stays flat rather than tracking DP size on top of it. + num_replicas = 4 + else: + # Each replica is a single event loop, so frontend capacity has to scale with + # the number of engines it feeds. The floor of 4 preserves the previous default + # for small deployments. + num_replicas = max(get_pg_size(engine.pg_collection.dp), 4) + if rank == 0: + logging.info("Starting %d HTTP frontend replica(s) per hosting rank.", num_replicas) + + if args.frontend_on_all_ranks: + # Only the DP coordinator rank learns the coordinator's address: the + # engine broadcasts it over the DP group, which is a singleton when data + # parallel size is 1. Every rank needs it here, since every rank's + # frontend opens its own client. + address = [coordinator_addr] + torch.distributed.broadcast_object_list(address, src=0) + coordinator_addr = address[0] + assert coordinator_addr is not None, "no rank published a DP coordinator address" + try: - if rank == 0: - start_text_gen_server( + url = None + if args.frontend_on_all_ranks or rank == 0: + url = start_text_gen_server( coordinator_addr=coordinator_addr, tokenizer=engine.controller.tokenizer, parsers=args.parsers, rank=rank, - server_port=server_port, + server_port=0 if args.frontend_on_all_ranks else server_port, verbose=args.inference_text_gen_server_logging, + num_replicas=num_replicas, hostname=hostname, ) + if args.frontend_on_all_ranks: + # Unlike callers that already collect a URL per worker, this entry + # point has to gather them itself before it can report the set. + urls = [None] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(urls, url) + if rank == 0: + for entry in [u for u in urls if u]: + logging.info("Frontend: %s", entry) + elif rank == 0: + logging.info("Frontend: %s", url) + # Await the engine loop directly since the server is running in a separate process await engine.engine_loop_task finally: - # Guarantee that the separate process is terminated when the engine loop stops or is interrupted - if rank == 0: - stop_text_gen_server() + # Guarantee that the separate processes are terminated when the engine loop + # stops or is interrupted. Every rank may now own frontend processes. + stop_text_gen_server() if __name__ == "__main__":