diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index e66bafd817..10b5e88d29 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -721,7 +721,12 @@ class ElasticEPScalingExecutor: self._suppress_eplb(), serving_state, ): - runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True) + runner._dummy_run( + runner.max_num_tokens, + is_profile=True, + skip_eplb=True, + randomize_inputs=self.worker.randomize_dummy_inputs, + ) self.worker.compile_or_warm_up_model() lock_workspace() diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 14ff11f448..b40fb4eac1 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -805,6 +805,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): skip_eplb: bool = False, is_profile: bool = False, valid_dummy_state_slots: bool = False, + randomize_inputs: bool = False, **kwargs, ) -> tuple[torch.Tensor | None, torch.Tensor | None]: if skip_attn and not is_profile: @@ -878,6 +879,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): is_profile=is_profile, context_len=context_len, valid_dummy_state_slots=valid_dummy_state_slots, + randomize_inputs=randomize_inputs, ) self.kv_connector.set_disabled(False) @@ -968,7 +970,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.pooling_runner.dummy_pooler_run(hidden_states) @torch.inference_mode() - def profile_run(self) -> None: + def profile_run(self, randomize_inputs: bool = False) -> None: if self.supports_mm_inputs and self.is_first_pp_rank: mm_config = self.model_config.multimodal_config if mm_config is not None and not mm_config.skip_mm_profiling: @@ -981,7 +983,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): ) hidden_states, sample_hidden_states = self._dummy_run( - self.max_num_tokens, skip_attn=True, is_profile=True + self.max_num_tokens, + skip_attn=True, + is_profile=True, + randomize_inputs=randomize_inputs, ) # Only run sampler/pooler on last PP rank (non-last ranks return None). @@ -1696,6 +1701,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): is_profile: bool = False, context_len: int = 0, valid_dummy_state_slots: bool = False, + randomize_inputs: bool = False, ) -> ModelRunnerOutput | IntermediateTensors | None: if not dummy_run: # Update the request states. @@ -1814,6 +1820,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): # so MoE memory is measured and MoE kernels are exercised. is_padding=not is_profile, ) + if randomize_inputs: + # All-zero input_ids route every token to the same experts. + input_batch.input_ids.random_(0, self.vocab_size) if self.pcp_manager is not None: input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch) if skip_attn_for_dummy_run: diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 3dc0524bab..4ccdec2f21 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5735,6 +5735,9 @@ class GPUModelRunner( - during profile_run - during DP rank dummy run """ + # The worker also resolves this flag and passes randomize_inputs, but + # V1 has internal dummy runs (e.g. CUDA graph capture) that don't go + # through the worker, so keep the check here too. dp_size = self.vllm_config.parallel_config.data_parallel_size randomize_inputs = randomize_inputs or ( envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS and dp_size > 1 @@ -6420,7 +6423,7 @@ class GPUModelRunner( max_task = max(output_size.items(), key=lambda x: x[1])[0] return self._dummy_pooler_run_task(hidden_states, max_task) - def profile_run(self) -> None: + def profile_run(self, randomize_inputs: bool = False) -> None: # Profile with multimodal encoder & encoder cache. if self.supports_mm_inputs: mm_config = self.model_config.multimodal_config @@ -6481,7 +6484,7 @@ class GPUModelRunner( # Add `is_profile` here to pre-allocate communication buffers hidden_states, last_hidden_states = self._dummy_run( - self.max_num_tokens, is_profile=True + self.max_num_tokens, is_profile=True, randomize_inputs=randomize_inputs ) if get_pp_group().is_last_rank: if self.is_pooling_model: diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 247265b659..e9021bd142 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -613,7 +613,7 @@ class Worker(WorkerBase): if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes: # still need a profile run which compiles the model for # max_num_batched_tokens - self.model_runner.profile_run() + self.model_runner.profile_run(randomize_inputs=self.randomize_dummy_inputs) msg = ( f"Initial free memory {format_gib(self.init_snapshot.free_memory)} " @@ -650,7 +650,7 @@ class Worker(WorkerBase): # memory_profiling measures, restoring the original limit. self._scoped_allocator_max_split(max_split_size_mb=20), ): - self.model_runner.profile_run() + self.model_runner.profile_run(randomize_inputs=self.randomize_dummy_inputs) # Profile CUDA graph memory if graphs will be captured. # ROCm is included: #44825 moved the profiler to @@ -885,7 +885,12 @@ class Worker(WorkerBase): # We skip EPLB here since we don't want to record dummy metrics for size in sorted(warmup_sizes, reverse=True): logger.info("Compile and warming up model for size %d", size) - self.model_runner._dummy_run(size, skip_eplb=True, remove_lora=False) + self.model_runner._dummy_run( + size, + skip_eplb=True, + remove_lora=False, + randomize_inputs=self.randomize_dummy_inputs, + ) self.model_runner.maybe_remove_all_loras(self.model_runner.lora_config) # Warmup and tune the kernels used during model execution before @@ -1387,9 +1392,21 @@ class Worker(WorkerBase): # Recreate it so the next profile_prefix is honored. self.profiler = None + @property + def randomize_dummy_inputs(self) -> bool: + # Not cached: elastic EP rewrites parallel_config in place on reconfigure. + return ( + envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS + and self.parallel_config.data_parallel_size > 1 + ) + def execute_dummy_batch(self) -> None: num_tokens = getattr(self.model_runner, "uniform_decode_query_len", 1) - self.model_runner._dummy_run(num_tokens, uniform_decode=True) + self.model_runner._dummy_run( + num_tokens, + uniform_decode=True, + randomize_inputs=self.randomize_dummy_inputs, + ) def add_lora(self, lora_request: LoRARequest) -> bool: return self.model_runner.add_lora(lora_request)