chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[CPU] support the case where num_attention_heads or intermediate_size is not divisible by the TP size (#6771)

Browse Source
This commit is contained in:
Chunyuan WU
2025-07-04 00:51:38 +08:00
committed by GitHub
parent 9fcc9a80e7
commit 1dce6c480f
11 changed files with 399 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

91
python/sglang/srt/layers/linear.py
View File

@@ -426,8 +426,26 @@ class ColumnParallelLinear(LinearBase):
if output_dim is not None and not use_bitsandbytes_4bit:
shard_size = param_data.shape[output_dim]
start_idx = self.tp_rank * shard_size
if not self.use_presharded_weights:
loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
if _is_cpu:
from sglang.srt.model_loader.weight_utils import (
narrow_padded_param_and_loaded_weight,
)
param_data, loaded_weight = narrow_padded_param_and_loaded_weight(
param_data,
loaded_weight,
0, # param_data_start
start_idx,
output_dim,
shard_size,
not self.use_presharded_weights,
)
else:
if not self.use_presharded_weights:
loaded_weight = loaded_weight.narrow(
output_dim, start_idx, shard_size
)
# Special case for loading scales off disk, which often do not
# have a shape (such as in the case of AutoFP8).
@@ -644,10 +662,29 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
param_data = param_data.narrow(output_dim, shard_offset, shard_size)
start_idx = self.tp_rank * shard_size
# bitsandbytes loads the weights of the specific portion
# no need to narrow here
if not use_bitsandbytes_4bit and not self.use_presharded_weights:
loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
if _is_cpu:
from sglang.srt.model_loader.weight_utils import (
narrow_padded_param_and_loaded_weight,
)
param_data, loaded_weight = narrow_padded_param_and_loaded_weight(
param_data,
loaded_weight,
0, # param_data_start
start_idx,
output_dim,
shard_size,
not use_bitsandbytes_4bit and not self.use_presharded_weights,
)
else:
# bitsandbytes loads the weights of the specific portion
# no need to narrow here
if not use_bitsandbytes_4bit and not self.use_presharded_weights:
loaded_weight = loaded_weight.narrow(
output_dim, start_idx, shard_size
)
# Special case for AQLM codebooks.
elif is_metadata:
# metadata indicates fixed size concatenated along dim 0
@@ -1112,10 +1149,27 @@ class QKVParallelLinear(ColumnParallelLinear):
shard_id = self.tp_rank // self.num_kv_head_replicas
start_idx = shard_id * shard_size
# bitsandbytes loads the weights of the specific portion
# no need to narrow here
if not use_bitsandbytes_4bit and not self.use_presharded_weights:
loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
if _is_cpu:
from sglang.srt.model_loader.weight_utils import (
narrow_padded_param_and_loaded_weight,
)
param_data, loaded_weight = narrow_padded_param_and_loaded_weight(
param_data,
loaded_weight,
0, # param_data_start
start_idx,
output_dim,
shard_size,
not use_bitsandbytes_4bit and not self.use_presharded_weights,
)
else:
# bitsandbytes loads the weights of the specific portion
# no need to narrow here
if not use_bitsandbytes_4bit and not self.use_presharded_weights:
loaded_weight = loaded_weight.narrow(
output_dim, start_idx, shard_size
)
# Special case for for AQLM codebooks.
elif is_metadata:
@@ -1257,7 +1311,22 @@ class RowParallelLinear(LinearBase):
):
shard_size = param_data.shape[input_dim]
start_idx = self.tp_rank * shard_size
loaded_weight = loaded_weight.narrow(input_dim, start_idx, shard_size)
if _is_cpu:
from sglang.srt.model_loader.weight_utils import (
narrow_padded_param_and_loaded_weight,
)
param_data, loaded_weight = narrow_padded_param_and_loaded_weight(
param_data,
loaded_weight,
0, # param_data_start
start_idx,
input_dim,
shard_size,
)
else:
loaded_weight = loaded_weight.narrow(input_dim, start_idx, shard_size)
# Special case for loading scales off disk, which often do not
# have a shape (such as in the case of AutoFP8).