Merge pull request #203 from Anai-Guo/fix-row-parallel-bias-crash

fix RowParallelLinear weight_loader crash when bias is enabled
2026-04-13 21:26:07 +08:00
parent a4f94cb38b bf99453d90
commit 02a95fdc66
1 changed files with 4 additions and 0 deletions
@@ -141,6 +141,10 @@ class RowParallelLinear(LinearBase):

    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
        param_data = param.data
+        if param_data.dim() == 1:
+            # bias is not sharded in RowParallelLinear
+            param_data.copy_(loaded_weight)
+            return
        shard_size = param_data.size(self.tp_dim)
        start_idx = self.tp_rank * shard_size
        loaded_weight = loaded_weight.narrow(self.tp_dim, start_idx, shard_size)