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)