Rain-Bus/nano-vllm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
fix RowParallelLinear weight_loader crash when bias is enabled
This commit is contained in:
Xingkai Yu
2026-04-13 21:26:07 +08:00
committed by GitHub
parent a4f94cb38b bf99453d90
commit 02a95fdc66
1 changed files with 4 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanovllm/layers/linear.py
+4
View File
@@ -141,6 +141,10 @@ class RowParallelLinear(LinearBase):
def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
param_data = param.data 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) shard_size = param_data.size(self.tp_dim)
start_idx = self.tp_rank * shard_size start_idx = self.tp_rank * shard_size
loaded_weight = loaded_weight.narrow(self.tp_dim, start_idx, shard_size) loaded_weight = loaded_weight.narrow(self.tp_dim, start_idx, shard_size)