diff --git a/nemo_deploy/llm/inference/inference_base.py b/nemo_deploy/llm/inference/inference_base.py index 129344179..6d72ef948 100644 --- a/nemo_deploy/llm/inference/inference_base.py +++ b/nemo_deploy/llm/inference/inference_base.py @@ -536,7 +536,11 @@ def create_mcore_engine( buffer_size_gb=int(buffer_size_gb), max_requests=max_batch_size, block_size_tokens=block_size_tokens, - materialize_only_last_token_logits=True, + # Coordinator mode constructs one DynamicInferenceEngine per rank independently, so a + # per-request runtime toggle (as _infer_fn used to do) only ever reaches the primary + # rank's engine. Prompt log probs require this to be False on every rank, so it's set + # here unconditionally at construction time instead. + materialize_only_last_token_logits=False, ) coordinator_host = os.environ.get("MASTER_ADDR") or _default_coordinator_host() diff --git a/nemo_deploy/llm/megatronllm_deployable.py b/nemo_deploy/llm/megatronllm_deployable.py index bec928698..0fc2d003b 100755 --- a/nemo_deploy/llm/megatronllm_deployable.py +++ b/nemo_deploy/llm/megatronllm_deployable.py @@ -205,12 +205,6 @@ def generate_other_ranks(self): stop_words=stop_words, ) - if log_probs: - dynamic_engine = getattr(self.mcore_engine, "engine", None) - if dynamic_engine is not None: - dynamic_engine.materialize_only_last_token_logits = False - dynamic_engine.context.config.materialize_only_last_token_logits = False - self.generate(prompts, inference_params) else: return @@ -400,22 +394,6 @@ def _infer_fn( if apply_chat_template: prompts = [self.apply_chat_template(prompt) for prompt in prompts] - if torch.distributed.is_initialized(): - if torch.distributed.get_world_size() > 1: - torch.distributed.broadcast(torch.tensor([0], dtype=torch.long, device="cuda"), src=0) - broadcast_list(prompts, src=0) - broadcast_list( - data=[ - temperature, - top_k, - top_p, - num_tokens_to_generate, - log_probs, - stop_words, - ], - src=0, - ) - # cast top_k,top_p to native int, float since typecheck assert statements added in MCore0.13 error otherwise # skip_prompt_log_probs=False (default) includes prompt tokens in top-N logprobs when top_logprobs>0. inference_params = SamplingParams( @@ -429,17 +407,6 @@ def _infer_fn( stop_words=stop_words, ) - # Mcore's dynamic inference engine defaults materialize_only_last_token_logits=True for - # performance, but prompt log probs require all token logits to be materialized - # (prompt log probs are required for logprob eval benchmarks). - # Toggle it on both the engine and the context config (controls the - # model forward pass and log prob calculations). - dynamic_engine = getattr(self.mcore_engine, "engine", None) - needs_all_logits = log_probs or bool(top_logprobs) - if dynamic_engine is not None and needs_all_logits: - dynamic_engine.materialize_only_last_token_logits = False - dynamic_engine.context.config.materialize_only_last_token_logits = False - results = self.generate(prompts, inference_params) # Handle DynamicInferenceRequestRecord objects by merging them into a single request results = [ diff --git a/nemo_deploy/llm/megatronllm_deployable_ray.py b/nemo_deploy/llm/megatronllm_deployable_ray.py index 3d2202aa9..8e4d326c4 100644 --- a/nemo_deploy/llm/megatronllm_deployable_ray.py +++ b/nemo_deploy/llm/megatronllm_deployable_ray.py @@ -104,8 +104,6 @@ def __init__( legacy_model_format=legacy_model_format, **model_config_kwargs, ) - if rank != 0: - self.model.generate_other_ranks() except Exception as e: LOGGER.error(f"Replica {replica_id} - Failed to initialize model for rank {rank}: {str(e)}") raise