【bug】RowParallelLinear adds bias multiple times when tp_size > 1

[HaoYuan-Gao]

I found a possible bug in RowParallelLinear.forward(). Current implementation: result = nn.functional.linear(x, self.weight, self.bias)

if self.tp_size > 1:
    dist.all_reduce(result, op=dist.ReduceOp.SUM)
return result

For RowParallelLinear, each rank computes a partial output and then uses all_reduce(SUM) to get the final result. However, because bias is added before all_reduce, each rank adds the same bias once. After all_reduce, the final result contains tp_size * bias. For example, when tp_size = 2, the output has one extra bias.

Test script: test.py
Run: PYTHONPATH=$PWD/src torchrun --nproc_per_node=2 test.py

Output:

Reference full linear output:
tensor([[130., 270., 410.],
        [170., 374., 578.]])

RowParallelLinear output:
tensor([[230., 470., 710.],
        [270., 574., 878.]], grad_fn=<AddmmBackward0>)

Difference:
tensor([[100., 200., 300.],
        [100., 200., 300.]], grad_fn=<SubBackward0>)

Expected extra bias:
tensor([100., 200., 300.])

Max error:
300.0

The difference is exactly one extra bias.

Additional Notes

The current loader design attaches the same weight_loader to both weight and bias:

self.weight.weight_loader = self.weight_loader 
self.bias.weight_loader = self.weight_loader

This works for ColumnParallelLinear, because both weight and bias are sharded along dim=0:

weight: [out_features, in_features] -> shard dim=0 
bias: [out_features] -> shard dim=0

However, this does not work for RowParallelLinear.