[Bug] Fix AttributeError: 'ColumnParallelLinear' object has no attribute weight_scale_inv

[yewentao256]

Purpose

Fix AttributeError: 'ColumnParallelLinear' object has no attribute weight_scale_inv

export MODEL="RedHatAI/Kimi-K2-Thinking-FP8-Block"

vllm serve $MODEL -tp 8 --port 9256 --enable-expert-parallel --enforce_eager --trust_remote_code --gpu_memory_utilization 0.94 --max_model_len 4096

(Worker_TP0_EP0 pid=1089072) INFO 12-16 23:53:39 [default_loader.py:308] Loading weights took 36.23 seconds
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751] WorkerProc failed to start.
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751] Traceback (most recent call last):
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]   File "/home/yewentao256/vllm-source/vllm/v1/executor/multiproc_executor.py", line 722, in worker_main
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]     worker = WorkerProc(*args, **kwargs)
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]   File "/home/yewentao256/vllm-source/vllm/v1/executor/multiproc_executor.py", line 562, in __init__
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]     self.worker.load_model()
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]   File "/home/yewentao256/vllm-source/vllm/v1/worker/gpu_worker.py", line 289, in load_model
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]     self.model_runner.load_model(eep_scale_up=eep_scale_up)
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]   File "/home/yewentao256/vllm-source/vllm/v1/worker/gpu_model_runner.py", line 3588, in load_model
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]     self.model = model_loader.load_model(
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]                  ^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]   File "/home/yewentao256/vllm-source/vllm/model_executor/model_loader/base_loader.py", line 56, in load_model
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]     process_weights_after_loading(model, model_config, target_device)
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]   File "/home/yewentao256/vllm-source/vllm/model_executor/model_loader/utils.py", line 108, in process_weights_after_loading
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]     quant_method.process_weights_after_loading(module)
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]   File "/home/yewentao256/vllm-source/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py", line 896, in process_weights_after_loading
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]     layer.scheme.process_weights_after_loading(layer)
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]   File "/home/yewentao256/vllm-source/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py", line 176, in process_weights_after_loading
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]     maybe_post_process_fp8_weight_block(layer)
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]   File "/home/yewentao256/vllm-source/vllm/model_executor/layers/quantization/utils/fp8_utils.py", line 1442, in maybe_post_process_fp8_weight_block
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]     ws=layer.weight_scale_inv.data,
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]        ^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]   File "/home/yewentao256/.venv/lib64/python3.12/site-packages/torch/nn/modules/module.py", line 1964, in __getattr__
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751]     raise AttributeError(
(Worker_TP0_EP0 pid=1089072) ERROR 12-16 23:53:40 [multiproc_executor.py:751] AttributeError: 'ColumnParallelLinear' object has no attribute 'weight_scale_inv'. Did you mean: 'weight_scale'?

Now it can run succesfully

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[gemini-code-assist]

Code Review

This pull request effectively resolves an AttributeError that occurs during the loading of certain FP8 models. The issue stemmed from hardcoding the weight_scale_inv attribute, which is not always present. The fix introduces a dynamic check for weight_scale_inv and falls back to weight_scale if the former is not found. This approach is robust and correctly handles different naming conventions for weight scales. The change is well-implemented and directly addresses the bug. I approve this pull request.