Update vllm/model_executor/model_loader/loader.py

Co-authored-by: Simon Mo <simon.mo@hey.com> Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
ROCm · Jan 31, 2025 · 1621381 · 1621381
1 parent 9829fae
commit 1621381
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vllm/model_executor/model_loader/loader.py b/vllm/model_executor/model_loader/loader.py
@@ -644,7 +644,7 @@ def load_model(self, vllm_config: VllmConfig) -> nn.Module:
                         hasattr(module, "process_weights_after_loading"):
                         # When attention modules need to process weights after
                         # currently only used by MLA
-                        module.process_weights_after_loading()
+                        module.process_weights_after_loading(model_config.dtype)
             rank = get_tensor_model_parallel_rank()
             pattern = os.path.join(
                 local_model_path,