Repository navigation
How to achieve vLLM Row Parallelism correctly? #186
Description
I am trying to make bitblas quantized weights to work with vLLM's tensor parallelism. In vLLM, tensor parallelism is achieved with column parallelism and row parallelism.
Similar to the reference vLLM integration here qweight, scales and zeros have the following shapes, input dims and output dims:
qweight = Parameter(
torch.empty(
output_size_per_partition,
input_size_per_partition // self.layer_pack_factor,
device="cuda",
dtype=torch.uint8,
),
requires_grad=False,
)
set_weight_attrs(
qweight,
{
"input_dim": 1,
"output_dim": 0,
"packed_dim": 1,
"pack_factor": self.layer_pack_factor,
},
)
scales = Parameter(
torch.empty(
output_size_per_partition,
input_size_per_partition // self.layer_group_size,
device="cuda",
dtype=params_dtype,
),
requires_grad=False,
)
set_weight_attrs(
scales,
{
"input_dim": 1,
"output_dim": 0,
},
)
zeros = Parameter(
torch.empty(
output_size_per_partition,
input_size_per_partition // self.layer_group_size,
device="cuda",
dtype=params_dtype,
),
requires_grad=False,
)
set_weight_attrs(
zeros,
{
"input_dim": 1,
"output_dim": 0,
},
)Our vLLM bitblas integration can be found here.
Column Parallel
In column parallelism, parallelization is done along the weight's output dim and the outputs are concatenated along the output's output dim using an all gather operation. Testing this behavior with bitblas works as expected:
# Mimic the output without any tensor parallelism.
bitblas_output = matmul_eng(x, Wq_bitblas, scales_bitblas, zeros_bitblas)
# Reconfigure matmul eng for new dims:
matmul_config = bitblas.MatmulConfig(M=BITBLAS_OPT_M,
N=out_features//2,
K=in_features,
A_dtype="float16",
W_dtype={4:"uint4",2:"uint2"}[NBITS],
accum_dtype="float16",
out_dtype="float16",
layout="nt",
with_bias=False,
group_size=GROUPSIZE,
with_scaling=True,
with_zeros=True,
zeros_mode="original",
#fast_decoding=True,
)
matmul_eng = _get_or_create_bitblas_operator(matmul_config)
# Split weights along the input dimenstion which is the 'N' dimension from matmul eq: (M x K @ K x N).
Wq_bitblas_split_1, Wq_bitblas_split_2 = Wq_bitblas.split(split_size=Wq_bitblas.size(0) // 2, dim=0)
zeros_bitblas_split_1, zeros_bitblas_split_2 = zeros_bitblas.split(split_size=zeros_bitblas.size(0) // 2, dim=0)
scales_bitblas_split_1, scales_bitblas_split_2 = scales_bitblas.split(split_size=scales_bitblas.size(0) // 2, dim=0)
bitblas_output_split_1 = matmul_eng(x, Wq_bitblas_split_1, scales_bitblas_split_1, zeros_bitblas_split_1)
bitblas_output_split_2 = matmul_eng(x, Wq_bitblas_split_2, scales_bitblas_split_2, zeros_bitblas_split_2)
# Test passes.
bitblas_sharded_output = torch.cat([bitblas_output_split_1, bitblas_output_split_2], dim=1)
assert torch.allclose(bitblas_sharded_output, bitblas_output, atol=1e-2, rtol=1e-2)Row Parallel
In row parallelism, parallelization is done along weight's input dim (K) and reduced (summed) along the via all reduce (sum) operation.
The outputs from the following test is very different compared to without tensor parallel.
# Mimic the output without any tensor parallelism.
bitblas_output = matmul_eng(x, Wq_bitblas, scales_bitblas, zeros_bitblas)
# Reconfigure matmul eng for new dims:
matmul_config = bitblas.MatmulConfig(M=BITBLAS_OPT_M,
N=out_features,
K=in_features//2,
A_dtype="float16",
W_dtype={4:"uint4",2:"uint2"}[NBITS],
accum_dtype="float16",
out_dtype="float16",
layout="nt",
with_bias=False,
group_size=GROUPSIZE,
with_scaling=True,
with_zeros=True,
zeros_mode="original",
#fast_decoding=True,
)
matmul_eng = _get_or_create_bitblas_operator(matmul_config)
# Split weights along the input dimenstion which is the 'K' dimension from matmul eq: (M x K @ K x N).
Wq_bitblas_split_1, Wq_bitblas_split_2 = Wq_bitblas.split(split_size=Wq_bitblas.size(1) // 2, dim=1)
zeros_bitblas_split_1, zeros_bitblas_split_2 = zeros_bitblas.split(split_size=zeros_bitblas.size(1) // 2, dim=1)
scales_bitblas_split_1, scales_bitblas_split_2 = scales_bitblas.split(split_size=scales_bitblas.size(1) // 2, dim=1)
# Also split the input along the K dimension.
x_1, x_2 = x.split(split_size=x.size(1) // 2, dim=1)
bitblas_output_split_1 = matmul_eng(x_1, Wq_bitblas_split_1, scales_bitblas_split_1, zeros_bitblas_split_1)
bitblas_output_split_2 = matmul_eng(x_2, Wq_bitblas_split_2, scales_bitblas_split_2, zeros_bitblas_split_2)
# Test fails
bitblas_sharded_output = bitblas_output_split_1 + bitblas_output_split_2
assert torch.allclose(bitblas_sharded_output, bitblas_output, atol=1e-2, rtol=1e-2)
Is there a wrong assumption here regarding the layout and/or packing? Maybe the zeros and scales are not correctly split? Appreciate your help. Thanks.
Activity
LeiWang1999 commentedon Sep 19, 2024
Kerem Turgutlu (@KeremTurgutlu) ，Thanks for your reporting! I'll take a look soon
LeiWang1999 commentedon Sep 20, 2024
Kerem Turgutlu (@KeremTurgutlu) Do we have a unit test script for us to reproduce? I can reproduce with tp=2 on a end2end model.
KeremTurgutlu commentedon Sep 20, 2024
Yes I think tp=2 should be a good start. Easiest would be to rerun the python code samples I shared in this issues because they isolate the issue to a single matmul rather than end2end model testing.
KeremTurgutlu commentedon Sep 27, 2024
Lei Wang (@LeiWang1999) Anything else needed from my side? Thanks.
LeiWang1999 commentedon Sep 27, 2024
Zhengju Tang (@tzj-fxz) thanks, I'm working on it, the last few days have been busy.