Sun Ruoxi
a53f6a02d0
All checks were successful
Docker Build and Push / docker (push) Successful in 1m4s
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
182 lines
7.9 KiB
Python
182 lines
7.9 KiB
Python
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)
|