add log in loading models
All checks were successful
Docker Build and Push / docker (push) Successful in 1m4s
All checks were successful
Docker Build and Push / docker (push) Successful in 1m4s
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
181
patch_vllm_load_progress.py
Normal file
181
patch_vllm_load_progress.py
Normal file
@@ -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)
|
||||
Reference in New Issue
Block a user