Fix ReplicatedLinear weight loading (vllm-project#6793)

OpenBMB · Jul 26, 2024 · 062a1d0 · 062a1d0
1 parent 2eb9f4f
commit 062a1d0
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py
@@ -199,12 +199,16 @@ def __init__(self,
                                          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)