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

fix RowParallelLinear weight_loader crash when bias is enabled

Browse Source 
When RowParallelLinear has bias=True, the weight_loader crashes with an
IndexError because it calls param_data.size(tp_dim) where tp_dim=1, but
the bias tensor is 1D and only has dimension 0.

The bias in RowParallelLinear is not sharded (all ranks hold the full
bias, only rank 0 applies it), so skip the sharding logic for 1D params.

Fixes GeeeekExplorer/nano-vllm#125
This commit is contained in:
Anai-Guo
2026-04-12 23:10:52 -07:00
parent 2f21442653
commit bf99453d90
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)