add log in loading models

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

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

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)