[Feature] Improve weight loading log (#18651)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Praneth Paruchuri
2026-03-04 03:46:13 +05:30
committed by GitHub
parent 9305f0e58d
commit f7897def96
3 changed files with 109 additions and 4 deletions

View File

@@ -775,14 +775,36 @@ class ModelConfig:
quant_algo = json_quant_configs.get("quant_algo", None)
if quant_algo == "MIXED_PRECISION":
return {"quant_method": "w4afp8"}
return {"quant_method": "w4afp8", "quant_algo": quant_algo}
elif quant_algo and ("FP4" in quant_algo or "NVFP4" in quant_algo):
return {"quant_method": "modelopt_fp4"}
return {"quant_method": "modelopt_fp4", "quant_algo": quant_algo}
elif quant_algo and "FP8" in quant_algo:
return {"quant_method": "modelopt_fp8"}
return {"quant_method": "modelopt_fp8", "quant_algo": quant_algo}
else:
return None
def get_quantization_config_log_str(self) -> Optional[str]:
"""
Get a concise string representation of the quantization config for logging.
Returns something like "quant=fp8, fmt=e4m3" or "quant=gptq, bits=4".
"""
try:
quant_cfg = self._parse_quant_hf_config()
if not quant_cfg:
return None
quant_method = quant_cfg.get("quant_method", "quantized")
log_str = f"quant={quant_method}"
# Append interesting fields if they exist
for field in ["bits", "quant_algo", "fmt"]:
if field in quant_cfg:
log_str += f", {field}={quant_cfg[field]}"
return log_str
except Exception:
return None
def _is_already_quantized(self) -> bool:
"""Check if the model is already quantized based on config files."""
# Check for quantization in hf_config (config.json)

View File

@@ -1041,11 +1041,15 @@ class ModelRunner(ModelRunnerKVCacheMixin):
after_avail_memory = get_available_gpu_memory(self.device, self.gpu_id)
self.weight_load_mem_usage = before_avail_memory - after_avail_memory
# Get quantization config from ModelConfig
# This handles both config.json (standard) and hf_quant_config.json (ModelOpt)
quant_str = self.model_config.get_quantization_config_log_str()
logger.info(
f"Load weight end. "
f"elapsed={time.perf_counter() - tic_total:.2f} s, "
f"type={type(self.model).__name__}, "
f"dtype={self.dtype}, "
f"{quant_str + ', ' if quant_str else ''}"
f"avail mem={after_avail_memory:.2f} GB, "
f"mem usage={self.weight_load_mem_usage:.2f} GB."
)