fix dbrx weight loader (#212)

Co-authored-by: Charlie Fu <Charlie.Fu@amd.com>
ROCm · Oct 1, 2024 · 2d7ab9e · 2d7ab9e
1 parent 956b831
commit 2d7ab9e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vllm/model_executor/models/dbrx.py b/vllm/model_executor/models/dbrx.py
@@ -392,7 +392,7 @@ def sample(
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
 
         expert_params_mapping = [(
-            "w13_weight" if weight_name in ["w1", "v1"] else "w2_weight",
+            "w13_" if weight_name in ["w1", "v1"] else "w2_",
             f"mlp.{weight_name}.",
         ) for weight_name in ["w1", "v1", "w2"]]
         params_dict = dict(self.named_parameters(remove_duplicate=False))