Fix ReplicatedLinear weight loading (#6793)

062a1d0f · QQSong · GitHub · 2eb9f4ff · 062a1d0f
Unverified Commit 062a1d0f authored Jul 25, 2024 by QQSong Committed by GitHub Jul 25, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 1 deletion

vllm/model_executor/layers/linear.py vllm/model_executor/layers/linear.py +5 -1

No files found.
--- a/vllm/model_executor/layers/linear.py
+++ b/vllm/model_executor/layers/linear.py
@@ -199,12 +199,16 @@ class ReplicatedLinear(LinearBase):
                                         self.input_size,
                                         self.output_size,
                                         self.params_dtype,
+                                         weight_loader=self.weight_loader,
                                         prefix=prefix)

        if bias:
            self.bias = Parameter(
                torch.empty(self.output_size, dtype=self.params_dtype))
-            set_weight_attrs(self.bias, {"output_dim": 0})
+            set_weight_attrs(self.bias, {
+                "output_dim": 0,
+                "weight_loader": self.weight_loader,
+            })
        else:
            self.register_parameter("bias", None)