diff --git a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py index 247404d537e..aa42f492ca4 100644 --- a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py +++ b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py @@ -25,6 +25,7 @@ get_model_for_inference, ) from megatron.training import get_args, get_tokenizer, initialize_megatron +from megatron.core.utils import configure_nvtx_profiling # pylint: disable=line-too-long @@ -208,6 +209,7 @@ async def main( args_defaults={'no_load_rng': True, 'no_load_optim': True}, ) initialize_megatron() + configure_nvtx_profiling(True) tokenizer = get_tokenizer() diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py index 756f4fbe4be..df8e36c7bac 100644 --- a/megatron/core/inference/config.py +++ b/megatron/core/inference/config.py @@ -312,6 +312,16 @@ class InferenceConfig: performance variability for MoEs. """ + disable_ep_consensus: bool = False + """If True, the engine skips the EP-group consensus all-reduce in + `run_engine_with_coordinator` and decides whether to step based on local + state alone. The rank still calls `controller.dummy_forward()` whenever + `local_pending == 0`, so EP collectives (NCCL all-to-all, etc.) stay in + sync — without this, a peer running a real forward would deadlock waiting + on this rank's all-to-all participation. Trades off the consensus + all-reduce CPU cost for unconditional dummy_forwards on idle ranks. + """ + verbose: InitVar[bool] = False """Whether to log detailed context configuration at initialization. This is an InitVar and is not stored as a field on the config.""" diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index fcee2c1daef..e862c8dacd2 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -6,7 +6,6 @@ import math import multiprocessing import socket -import struct import time import warnings from collections import deque @@ -222,6 +221,7 @@ def __init__(self, controller: TextGenerationController, context: DynamicInferen self.logging_step_interval = inference_config.logging_step_interval self.unified_memory_level = inference_config.unified_memory_level self.use_synchronous_zmq_collectives = inference_config.use_synchronous_zmq_collectives + self.disable_ep_consensus = inference_config.disable_ep_consensus self.cuda_graph_impl = model_config.cuda_graph_impl self.cuda_graph_scope = model_config.cuda_graph_scope # Initialize engine. @@ -584,20 +584,16 @@ async def start_listening_to_data_parallel_coordinator( mp_req_sock.bind_to_random_port(f"tcp://{local_ip}") mp_req_addr = mp_req_sock.getsockopt_string(zmq.LAST_ENDPOINT) - mp_len_sock = self.zmq_context.socket(zmq.PUB) - mp_len_sock.bind_to_random_port(f"tcp://{local_ip}") - mp_len_addr = mp_len_sock.getsockopt_string(zmq.LAST_ENDPOINT) else: mp_req_addr = None - mp_len_addr = None # Broadcast addresses to respective ranks. bcast = [dp_addr] torch.distributed.broadcast_object_list(bcast, src=dp_src, group=dp_group) [dp_addr] = bcast - bcast = [mp_req_addr, mp_len_addr] + bcast = [mp_req_addr] torch.distributed.broadcast_object_list(bcast, src=mp_src, group=mp_group) - [mp_req_addr, mp_len_addr] = bcast + [mp_req_addr] = bcast identity = f'mp-coord-{dp_rank}' if self.is_mp_coordinator: @@ -614,33 +610,24 @@ async def start_listening_to_data_parallel_coordinator( # 2. Create a publisher socket. This is used to publish or broadcast # requests within the model parallel group self.model_parallel_publisher_socket = mp_req_sock - - # 3. Create another publisher socket to broadcast the number of messages to receive. - self.model_parallel_num_msgs_publisher_socket = mp_len_sock self.zmq_sockets += [ self.socket_for_receiving_requests, - self.model_parallel_num_msgs_publisher_socket, self.model_parallel_publisher_socket, ] - # All MP ranks subscribe to the two publisher sockets + # All MP ranks subscribe to the publisher socket self.model_parallel_subscriber_socket = self.zmq_context.socket(zmq.SUB) self.model_parallel_subscriber_socket.connect(mp_req_addr) self.model_parallel_subscriber_socket.setsockopt_string(zmq.SUBSCRIBE, "") - self.model_parallel_num_msgs_subscriber_socket = self.zmq_context.socket(zmq.SUB) - self.model_parallel_num_msgs_subscriber_socket.connect(mp_len_addr) - self.model_parallel_num_msgs_subscriber_socket.setsockopt_string(zmq.SUBSCRIBE, "") - - self.zmq_sockets += [ - self.model_parallel_subscriber_socket, - self.model_parallel_num_msgs_subscriber_socket, - ] + self.zmq_sockets += [self.model_parallel_subscriber_socket] torch.distributed.barrier(mp_group) # initialize zmq-based EP communicator self.ep_rank = get_pg_rank(self.pg_collection.ep) self.ep_world_size = get_pg_size(self.pg_collection.ep) + self._ep_consensus_loop_counter = 0 + self._last_ep_consensus: tuple[int, bool] = (0, False) if self.ep_world_size > 1: self.expert_parallel_zmq_communicator = AsyncZMQCommunicator( self.zmq_context, process_group=self.pg_collection.ep, hostname=hostname @@ -1870,6 +1857,7 @@ async def async_bookkeep( self.context.prefix_cache_blocks_matched = 0 # Log KV cache utilization stats to W&B + nvtx_range_push("wandb_logging") if context_state["kv_stats"] is not None: # Prepare metrics dictionary with all stats # Use 'inference/' prefix for all metrics to separate from training metrics @@ -1909,13 +1897,17 @@ async def async_bookkeep( self.metrics_writer.log(metrics, commit=True) else: raise ValueError(f"Unsupported metrics writer type: {type(self.metrics_writer)}") + nvtx_range_pop("wandb_logging") # Print context state. + nvtx_range_push("console_logging") if ( self.logging_step_interval > 0 and self.context.step_count % self.logging_step_interval == 0 ): + nvtx_range_push("cuda_memory_stats") mem = torch.cuda.memory_stats() + nvtx_range_pop("cuda_memory_stats") step_type = "decode" if context_state["is_decode_only"] else "non-decode" output_str = ( "* rank %d | step %d | %s ... time: %.3f ms%s ... " @@ -1982,6 +1974,8 @@ async def async_bookkeep( self._prefix_cache_hits = 0 self._prefix_cache_blocks_matched = 0 + nvtx_range_pop("console_logging") + return { "active_request_ids": active_request_ids, "finished_request_records": finished_request_records, @@ -2110,28 +2104,12 @@ def schedule_requests(self) -> int: except zmq.Again: # This exception is hit as soon as the socket is empty. break - messages_to_dequeue = len(all_messages) - # First publish the number of messages to dequeue. - # This is important because we want all tensor parallel ranks - # to dequeue the same number of messages. - self.model_parallel_num_msgs_publisher_socket.send( - struct.pack('!i', messages_to_dequeue) + self.model_parallel_publisher_socket.send_multipart( + [bytes([Headers.TP_BROADCAST.value])] + all_messages ) - # Now publish the actual messages to all model parallel ranks - if messages_to_dequeue > 0: - self.model_parallel_publisher_socket.send_multipart(all_messages) else: - # First, receive the number of messages to dequeue from mp-rank 0 - messages_to_dequeue = struct.unpack( - '!i', self.model_parallel_num_msgs_subscriber_socket.recv() - )[0] - # Now, dequeue the same number of messages from the subscriber socket. - # Note that these receives are blocking, because the messages - # are guaranteed to be available after the tp-rank 0 has sent them. - if messages_to_dequeue > 0: - all_messages = self.model_parallel_subscriber_socket.recv_multipart() - else: - all_messages = [] + frames = self.model_parallel_subscriber_socket.recv_multipart() + all_messages = frames[1:] nvtx_range_pop("drain_zmq_socket") @@ -2358,9 +2336,50 @@ async def run_engine_with_coordinator( local_pending = self.context.get_active_request_count() + len( self.waiting_request_ids ) - global_work, all_pausing = await self._ep_establish_consensus( - local_pending, signal_consensus=(self.state == EngineState.PAUSING) - ) + if self.disable_ep_consensus: + # Skip the EP consensus all-reduce; act on local state only. + # NOTE: even with no consensus we must still participate in EP + # collectives (NCCL all-to-all, etc.) every iteration. A peer with + # real work will block at its all-to-all kernel waiting for this + # rank, so when there is no local work we run dummy_forward() + # rather than sleeping. Sleeping here would deadlock EP > 1. + if self.state == EngineState.PAUSING: + await self._world_barrier() + self.state = EngineState.PAUSED + self._state_events[EngineState.PAUSED].set() + elif local_pending > 0: + await self.async_step() + else: + self.step_start_event.record() + nvtx_range_push("EP-dummy-forward") + self.controller.dummy_forward() + self.step_end_event.record() + self.step_end_event.synchronize() + nvtx_range_pop("EP-dummy-forward") + self.context.step_count += 1 + self.context.prefix_cache_lru_clock += 1 + # The consensus path yields via _ep_establish_consensus; + # without it we must still let other coroutines (signal + # delivery, request scheduling) run between steps. + await asyncio.sleep(0) + continue + global_work_from_last_consensus, _ = self._last_ep_consensus + if ( + global_work_from_last_consensus == 0 + or self._ep_consensus_loop_counter % 20 == 0 + ): + # selectively enter ep_establish_consensus if + # 1. there is no global work -> engine is idle. At any step in the future + # one of the ranks can receive work. So we should be eagerly checking for that + # 2. it has been 20 steps since we last established consensus, and that consensus + # had some work. + # In the worst case, this delays pausing by 20 steps which is around + # 200-400 milliseconds. + self._last_ep_consensus = await self._ep_establish_consensus( + local_pending, signal_consensus=(self.state == EngineState.PAUSING) + ) + global_work, all_pausing = self._last_ep_consensus + self._ep_consensus_loop_counter += 1 if all_pausing: # All EP peers are PAUSING: pause immediately. @@ -2374,9 +2393,11 @@ async def run_engine_with_coordinator( else: # Dummy forward to participate in the EP collective. self.step_start_event.record() + nvtx_range_push("EP-dummy-forward") self.controller.dummy_forward() self.step_end_event.record() self.step_end_event.synchronize() + nvtx_range_pop("EP-dummy-forward") self.context.step_count += 1 self.context.prefix_cache_lru_clock += 1 else: @@ -2391,6 +2412,10 @@ async def run_engine_with_coordinator( self.state = EngineState.RUNNING self._state_events[EngineState.PAUSED].clear() self._state_events[EngineState.RUNNING].set() + # The cache from the PAUSING phase still has all_pausing=True; + # without this reset the next RUNNING iteration would skip + # consensus, read the stale flag, and immediately re-pause. + self._last_ep_consensus = (0, False) elif self.state == EngineState.SUSPENDING: await self._world_barrier() diff --git a/megatron/core/inference/headers.py b/megatron/core/inference/headers.py index aa2f0568975..8ad1913e6b1 100644 --- a/megatron/core/inference/headers.py +++ b/megatron/core/inference/headers.py @@ -20,6 +20,7 @@ class Headers(Enum): STOP = auto() DISCONNECT = auto() SHUTDOWN = auto() + TP_BROADCAST = auto() class UnknownHeaderError(Exception): diff --git a/megatron/inference/utils.py b/megatron/inference/utils.py index 559ac57f496..0e14251c5aa 100644 --- a/megatron/inference/utils.py +++ b/megatron/inference/utils.py @@ -374,6 +374,7 @@ def get_inference_config_from_model_and_args(model: MegatronModule, args): logging_step_interval=args.inference_logging_step_interval, num_speculative_tokens=args.num_speculative_tokens, use_synchronous_zmq_collectives=args.inference_use_synchronous_zmq_collectives, + disable_ep_consensus=args.inference_disable_ep_consensus, sampling_backend=args.inference_dynamic_batching_sampling_backend, ) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index e79334b7c03..2040eb2bc7e 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2028,6 +2028,11 @@ def _add_inference_args(parser): help='Dtype for the Mamba inference SSM states tensor') group.add_argument('--inference-use-synchronous-zmq-collectives', action=argparse.BooleanOptionalAction, required=False, default=False, help='Use synchronous ZMQ collectives for inference. Helps in reducing performance variability for MoEs.') + group.add_argument('--inference-disable-ep-consensus', action=argparse.BooleanOptionalAction, + required=False, default=False, + help='Skip the EP-group consensus all-reduce in the inference engine control loop and step on local state only. ' + 'Pause/unpause take effect as soon as the signal is delivered to a rank. ' + 'Only safe when EP coordination is not required (e.g. ep_world_size == 1).') return parser diff --git a/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py b/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py index b2e94bc54f9..5f1aeca5b13 100644 --- a/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py +++ b/tests/unit_tests/inference/test_data_parallel_inference_coordinator.py @@ -118,6 +118,7 @@ def __init__(self): self.use_coordinator = False self.ep_world_size = 1 + self.disable_ep_consensus = False self.step_start_event = unittest.mock.MagicMock() self.step_end_event = unittest.mock.MagicMock() @@ -406,6 +407,69 @@ async def test_parallel_configs( finally: await cleanup_engine(engine, client) + @pytest.mark.internal + @pytest.mark.skipif(not HAVE_ZMQ, reason="pyzmq is required for this test") + @pytest.mark.asyncio + @pytest.mark.parametrize( + "initialize_model_parallel", + [pytest.param((1, 1, 1), id="tp1-pp1-ep1")], + indirect=["initialize_model_parallel"], + ) + async def test_disable_ep_consensus( + self, initialize_model_parallel, coordinator, test_case_communicator + ): + """With disable_ep_consensus=True, the control loop must call + controller.dummy_forward() on iterations where local_pending == 0 + instead of sleeping, so EP collectives stay in sync. Sleeping here + would deadlock peers running real forwards on EP > 1.""" + dp_addr = coordinator + port = int(dp_addr.rsplit(":", 1)[-1]) + requests = self.build_requests(num_requests=2) + engine = DummyEngine() + engine.disable_ep_consensus = True + engine.controller.dummy_forward = unittest.mock.MagicMock( + wraps=engine.controller.dummy_forward + ) + rank = torch.distributed.get_rank() + client = None + + try: + await engine.start_listening_to_data_parallel_coordinator( + inference_coordinator_port=port, launch_inference_coordinator=False + ) + await asyncio.wait_for(test_case_communicator.all_reduce_max(1), timeout=30.0) + + if rank == 0: + client = InferenceClient(dp_addr) + client.start() + await asyncio.wait_for(engine.wait_until(EngineState.RUNNING), timeout=5.0) + + # Idle window: with no work, the loop must spin on dummy_forward, + # not sleep. Several iterations should fire within 0.2s. + idle_baseline = engine.controller.dummy_forward.call_count + await asyncio.sleep(0.2) + idle_calls = engine.controller.dummy_forward.call_count - idle_baseline + assert idle_calls > 0, ( + "disable_ep_consensus must call dummy_forward on idle iterations " + f"to keep EP collectives in sync (call_count={idle_calls})" + ) + + # Submit and complete requests to confirm the step path still works. + futures = [client.add_request(prompt=p, sampling_params=s) for p, s in requests] + results = await asyncio.wait_for(asyncio.gather(*futures), timeout=5.0) + for result in results: + assert result["status"] == Status.COMPLETED.name + + # Pause/unpause must still drive state transitions correctly. + client.pause_engines() + await asyncio.wait_for(engine.wait_until(EngineState.PAUSED), timeout=5.0) + client.unpause_engines() + await asyncio.wait_for(engine.wait_until(EngineState.RUNNING), timeout=5.0) + + await asyncio.wait_for(test_case_communicator.all_reduce_max(1), timeout=30.0) + finally: + await cleanup_engine(engine, client) + @pytest.mark.internal @pytest.mark.skipif(not HAVE_ZMQ, reason="pyzmq is required for this test") @pytest.mark.asyncio