diff --git a/Dockerfile b/Dockerfile index a28975e..5bb51df 100644 --- a/Dockerfile +++ b/Dockerfile @@ -9,15 +9,23 @@ ENV JRE_HOME=/root/apps/jdk1.8.0_411/jre ENV JMETER_HOME=/root/apps/apache-jmeter-5.6.3 ENV CLASSPATH=.:/root/apps/jdk1.8.0_411/lib/dt.jar:/root/apps/jdk1.8.0_411/lib/tools.jar:/root/apps/apache-jmeter-5.6.3/lib/ext/ApacheJMeter_core.jar:/root/apps/apache-jmeter-5.6.3/lib/jorphan.jar:/root/apps/apache-jmeter-5.6.3/lib/logkit-2.0.jar: ENV PATH=/root/apps/apache-jmeter-5.6.3/bin:/root/apps/jdk1.8.0_411/bin:/usr/local/corex/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/local/corex/lib64/python3/dist-packages/bin:/usr/local/openmpi/bin +ENV PYTHONUNBUFFERED=1 +ENV VLLM_LOAD_PROGRESS_INTERVAL_SECONDS=60 +ENV VLLM_LOAD_PROGRESS_TENSOR_INTERVAL=50 COPY fix_tokenizer.py /opt/ COPY detect_tokenizer.py /opt/ COPY detect_head_size.py /opt/ COPY patch_ops.py /opt/ +COPY patch_vllm_load_progress.py /opt/ COPY patched_ops /opt/patched_ops/ COPY entrypoint.sh /opt/ # 在构建时执行 patch_ops（只需要执行一次） RUN python3 /opt/patch_ops.py && \ + python3 /opt/patch_vllm_load_progress.py && \ + python3 -m py_compile \ + /usr/local/corex/lib64/python3/dist-packages/vllm/model_executor/model_loader/loader.py \ + /usr/local/corex/lib64/python3/dist-packages/vllm/model_executor/model_loader/weight_utils.py && \ chmod +x /opt/entrypoint.sh diff --git a/patch_vllm_load_progress.py b/patch_vllm_load_progress.py new file mode 100644 index 0000000..7b6778e --- /dev/null +++ b/patch_vllm_load_progress.py @@ -0,0 +1,181 @@ +from pathlib import Path + + +SITE_PACKAGES = Path("/usr/local/corex/lib64/python3/dist-packages/vllm") +LOADER = SITE_PACKAGES / "model_executor/model_loader/loader.py" +WEIGHT_UTILS = SITE_PACKAGES / "model_executor/model_loader/weight_utils.py" + + +def replace_once(path: Path, old: str, new: str, description: str) -> None: + source = path.read_text() + count = source.count(old) + if count != 1: + raise RuntimeError( + f"Refusing to patch {description}: expected one match in {path}, got {count}" + ) + path.write_text(source.replace(old, new, 1)) + print(f"[vllm-load-progress-patch] patched {description}", flush=True) + + +replace_once( + WEIGHT_UTILS, + "import tempfile\n", + "import tempfile\nimport time\n", + "weight_utils time import", +) + +replace_once( + WEIGHT_UTILS, + ''' for st_file in tqdm( + hf_weights_files, + desc="Loading safetensors checkpoint shards", + disable=not enable_tqdm, + bar_format=_BAR_FORMAT, + ): + with safe_open(st_file, framework="pt") as f: + for name in f.keys(): # noqa: SIM118 + param = f.get_tensor(name) + yield name, param +''', + ''' progress_interval = max( + 1, int(os.getenv("VLLM_LOAD_PROGRESS_INTERVAL_SECONDS", "60"))) + tensor_interval = max( + 1, int(os.getenv("VLLM_LOAD_PROGRESS_TENSOR_INTERVAL", "50"))) + progress_started = time.monotonic() + last_progress_log = progress_started + loaded_tensors = 0 + loaded_bytes = 0 + total_files = len(hf_weights_files) + for file_index, st_file in enumerate(tqdm( + hf_weights_files, + desc="Loading safetensors checkpoint shards", + disable=not enable_tqdm, + bar_format=_BAR_FORMAT, + ), start=1): + if enable_tqdm: + logger.info( + "[VLLM_LOAD_PROGRESS] phase=load_weights status=file_start " + "file=%s file_index=%d total_files=%d", + os.path.basename(st_file), file_index, total_files) + with safe_open(st_file, framework="pt") as f: + for name in f.keys(): # noqa: SIM118 + param = f.get_tensor(name) + param_bytes = param.numel() * param.element_size() + yield name, param + # The generator resumes only after model.load_weights consumed + # this tensor, so this records completed rather than queued work. + loaded_tensors += 1 + loaded_bytes += param_bytes + now = time.monotonic() + if enable_tqdm and ( + loaded_tensors % tensor_interval == 0 + or now - last_progress_log >= progress_interval): + logger.info( + "[VLLM_LOAD_PROGRESS] phase=load_weights status=progress " + "file=%s file_index=%d total_files=%d tensors=%d " + "loaded_gib=%.2f elapsed_seconds=%.1f last_tensor=%s", + os.path.basename(st_file), file_index, total_files, + loaded_tensors, loaded_bytes / 1024**3, + now - progress_started, name) + last_progress_log = now + if enable_tqdm: + logger.info( + "[VLLM_LOAD_PROGRESS] phase=load_weights status=file_done " + "file=%s file_index=%d total_files=%d tensors=%d " + "loaded_gib=%.2f elapsed_seconds=%.1f", + os.path.basename(st_file), file_index, total_files, + loaded_tensors, loaded_bytes / 1024**3, + time.monotonic() - progress_started) +''', + "safetensors semantic progress", +) + +replace_once( + LOADER, + "import os\n", + "import os\nimport time\n", + "loader time import", +) + +replace_once( + LOADER, + ''' with set_default_torch_dtype(model_config.dtype): + with target_device: + model = _initialize_model(model_config, self.load_config, + lora_config, cache_config, + scheduler_config) + + model.load_weights(self._get_all_weights(model_config, model)) + + for _, module in model.named_modules(): + quant_method = getattr(module, "quant_method", None) + if quant_method is not None: + # When quant methods need to process weights after loading + # (for repacking, quantizing, etc), they expect parameters + # to be on the global target device. This scope is for the + # case where cpu offloading is used, where we will move the + # parameters onto device for processing and back off after. + with device_loading_context(module, target_device): + quant_method.process_weights_after_loading(module) + return model.eval() +''', + ''' with set_default_torch_dtype(model_config.dtype): + phase_started = time.monotonic() + logger.info( + "[VLLM_LOAD_PROGRESS] phase=initialize_model status=start model=%s", + model_config.model) + with target_device: + model = _initialize_model(model_config, self.load_config, + lora_config, cache_config, + scheduler_config) + logger.info( + "[VLLM_LOAD_PROGRESS] phase=initialize_model status=done " + "elapsed_seconds=%.1f", time.monotonic() - phase_started) + + phase_started = time.monotonic() + logger.info("[VLLM_LOAD_PROGRESS] phase=load_weights status=start") + model.load_weights(self._get_all_weights(model_config, model)) + logger.info( + "[VLLM_LOAD_PROGRESS] phase=load_weights status=done " + "elapsed_seconds=%.1f", time.monotonic() - phase_started) + + quant_modules = [ + (name, module, getattr(module, "quant_method", None)) + for name, module in model.named_modules() + if getattr(module, "quant_method", None) is not None + ] + phase_started = time.monotonic() + last_progress_log = phase_started + progress_interval = max( + 1, int(os.getenv("VLLM_LOAD_PROGRESS_INTERVAL_SECONDS", "60"))) + logger.info( + "[VLLM_LOAD_PROGRESS] phase=post_process_weights status=start " + "total_modules=%d", len(quant_modules)) + for module_index, (name, module, quant_method) in enumerate( + quant_modules, start=1): + # When quant methods need to process weights after loading + # (for repacking, quantizing, etc), they expect parameters + # to be on the global target device. + with device_loading_context(module, target_device): + quant_method.process_weights_after_loading(module) + now = time.monotonic() + if (now - last_progress_log >= progress_interval + or module_index == len(quant_modules)): + logger.info( + "[VLLM_LOAD_PROGRESS] phase=post_process_weights " + "status=progress modules=%d total_modules=%d " + "elapsed_seconds=%.1f last_module=%s", + module_index, len(quant_modules), + now - phase_started, name) + last_progress_log = now + logger.info( + "[VLLM_LOAD_PROGRESS] phase=post_process_weights status=done " + "total_modules=%d elapsed_seconds=%.1f", + len(quant_modules), time.monotonic() - phase_started) + logger.info("[VLLM_LOAD_PROGRESS] phase=load_model status=done") + return model.eval() +''', + "default model loader phase and quantization progress", +) + +print("[vllm-load-progress-patch] all patches applied", flush=True)