diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/llama-3_1-8B_pruneffn_runtime.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/llama-3_1-8B_pruneffn_runtime.yaml index 2ca0f2c16cf..e788443defd 100644 --- a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/llama-3_1-8B_pruneffn_runtime.yaml +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/llama-3_1-8B_pruneffn_runtime.yaml @@ -16,6 +16,12 @@ calc_subblock_stats: enabled: true num_warmup_iters: 2 num_iters: 10 + data_parallel_size: 1 + tensor_parallel_size: 1 + prefill_context_parallel_size: 1 + decode_context_parallel_size: 1 + enable_expert_parallel: false + extra_args: "" # FFN intermediate sizes to search over (heterogeneous architecture) pruning: diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py b/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py index c7d926de18f..669b25e675d 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py +++ b/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py @@ -186,9 +186,15 @@ def calc_runtime_for_subblocks( prefill_seq_len, generation_seq_len, batch_size, + runtime_stats_config.get("data_parallel_size", 1), + runtime_stats_config.get("tensor_parallel_size", 1), + runtime_stats_config.get("prefill_context_parallel_size", 1), + runtime_stats_config.get("decode_context_parallel_size", 1), + runtime_stats_config.get("enable_expert_parallel", False), runtime_stats_config.get("num_iters", 30), runtime_stats_config.get("num_warmup_iters", 10), runtime_stats_config.get("gpu_memory_utilization", 0.5), + runtime_stats_config.get("extra_args", ""), ) runtime_by_subblock_dict = {} diff --git a/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py b/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py index 204c2a74305..c25083692b0 100644 --- a/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py +++ b/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py @@ -50,6 +50,11 @@ class RuntimeConfig: prefill_seq_len: int generation_seq_len: int batch_size: int + data_parallel_size: int + tensor_parallel_size: int + prefill_context_parallel_size: int + decode_context_parallel_size: int + enable_expert_parallel: bool num_iters: int num_warmup_iters: int # Fraction of total GPU memory vLLM may use. Kept well below the default @@ -57,6 +62,7 @@ class RuntimeConfig: # GPU during benchmarking; requesting too much fails vLLM's startup # free-memory check. gpu_memory_utilization: float = 0.5 + extra_args: str = "" def save_model(model: LlamaForCausalLM, tokenizer_path: Path, output_path: Path) -> None: diff --git a/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py b/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py index 386f3615d49..bd2493d6115 100644 --- a/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py +++ b/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py @@ -27,6 +27,7 @@ """ import json +import os import subprocess # nosec B404 from pathlib import Path from types import SimpleNamespace @@ -79,12 +80,19 @@ def run_vllm_latency_benchmark(model_path: Path, runtime_config: RuntimeConfig) str(runtime_config.num_iters), "--max-num-seqs", "1", + "--data-parallel-size", + str(runtime_config.data_parallel_size), + "--prefill-context-parallel-size", + str(runtime_config.prefill_context_parallel_size), + "--decode-context-parallel-size", + str(runtime_config.decode_context_parallel_size), + "--enable-expert-parallel" + if runtime_config.enable_expert_parallel + else "--no-enable-expert-parallel", "--tensor-parallel-size", - "1", - "--pipeline-parallel-size", - "1", + str(runtime_config.tensor_parallel_size), "--distributed-executor-backend", - "external_launcher", + "mp", # "external_launcher", # Cap GPU memory so vLLM's startup free-memory check passes while the # parent puzzletron process is co-resident on the same GPU. "--gpu-memory-utilization", @@ -93,6 +101,26 @@ def run_vllm_latency_benchmark(model_path: Path, runtime_config: RuntimeConfig) "--optimization-level", "0", ] + if runtime_config.extra_args: + cmd.extend(runtime_config.extra_args.split(" ")) + + # Strip torchrun-injected distributed vars so vLLM's mp workers don't + # inherit an outer process group's rendezvous state, which causes a + # deadlock when vLLM tries to create its own Gloo CPU group. + _TORCHRUN_VARS = { + "RANK", + "LOCAL_RANK", + "WORLD_SIZE", + "LOCAL_WORLD_SIZE", + "MASTER_ADDR", + "MASTER_PORT", + "TORCHELASTIC_RESTART_COUNT", + "TORCHELASTIC_MAX_RESTARTS", + "TORCHELASTIC_RUN_ID", + "TORCHELASTIC_USE_AGENT_STORE", + "TORCHELASTIC_ERROR_FILE", + } + env = {k: v for k, v in os.environ.items() if k not in _TORCHRUN_VARS} # cmd is a fixed list of strings (no shell, no untrusted input). try: @@ -101,7 +129,8 @@ def run_vllm_latency_benchmark(model_path: Path, runtime_config: RuntimeConfig) check=True, capture_output=True, text=True, - timeout=1800, # 30 minutes + timeout=3600, # 1 hour + env=env, ) # nosec B603 except subprocess.TimeoutExpired as exc: raise TimeoutError("vLLM latency benchmark timed out") from exc