Fix ReplicatedLinear weight loading (#6793)

2024-07-25 19:24:58 -07:00 · 2024-07-25 19:24:58 -07:00 · 062a1d0fab
commit 062a1d0fab
parent 2eb9f4ff26
1 changed files with 5 additions and 1 deletions
--- 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)