A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 22:15 UTC
Correction to TP logic for Mamba Mixer 2 when Num Groups not divisible by TP Size (#13660)
Yu Chin Fabian Lim committed
Feb 22, 2025 at 16:19 UTC
fca20841c2467e2c811eebb004a0e6479edbe431
1 file changed
+21
-7
vllm/model_executor/layers/mamba/mamba_mixer2.py
+21
-7
@@ -133,7 +133,8 @@ def extra_groups_for_head_shards(ngroups: int, tp_size: int):
133
if ngroups % tp_size == 0:
134
return 0
135
136
- return tp_size - ngroups % tp_size
136
+ # for n_groups == 1, this is exactly tp_size - n_groups
137
+ return tp_size - ngroups
138
139
140
def mamba_v2_sharded_weight_loader(
@@ -153,7 +154,7 @@ def mamba_v2_sharded_weight_loader(
154
boundary, loaded_boundary = 0, 0
155
156
# - iterate over the shard specs
156
- for full_dim, extra, ratio in shard_spec:
157
+ for full_dim, extra, duplicate_groups in shard_spec:
158
# - full dim is the model dim (before TP).
159
# - extra > 0, means there is expected overall increase
160
# of dimensions. This is so because of replication.
@@ -167,7 +168,12 @@ def mamba_v2_sharded_weight_loader(
168
# - compute the rank into the loaded shard.
169
# - if there is replication, different TP shards will
170
# take from the same rank.
170
- rank = tp_rank // ratio
171
+ if duplicate_groups:
172
+ # NOTE: currently we only support duplication
173
+ # in the case where num_groups == 1
174
+ rank = 0
175
+ else:
176
+ rank = tp_rank
177
178
# - leftmost boundary index into loaded weight.
179
loaded_skip = rank * shard_size
@@ -233,12 +239,21 @@ class MambaMixer2(CustomOp):
239
# - HOWEVER IF, world_size DOES NOT divide groups, then we need
240
# to allocate extra space in the shard, such that groups
241
# may be replicated to follow the head shard.
242
+ # - NOTE: currently for the world size DOES NOT divide groups
243
+ # case, we only support the case when n_groups == 1
244
self.tp_size = get_tensor_model_parallel_world_size()
245
tp_rank = get_tensor_model_parallel_rank()
246
247
assert num_heads % self.tp_size == 0, \
248
"Tensor parallel world size must divide num heads."
249
250
+
251
+ assert (n_groups % self.tp_size) == 0 or n_groups == 1, \
252
+ (
253
+ "If tensor parallel world size does not divide num_heads, "
254
+ "then num_groups must equal 1."
255
+ )
256
+
257
self.ssm_state_size = ssm_state_size
258
self.activation = activation
259
@@ -284,11 +299,10 @@ class MambaMixer2(CustomOp):
299
self.n_groups * self.ssm_state_size, # expected model size
300
(self.n_groups - n_groups) *
301
self.ssm_state_size, # extra dims assigned
287
- self.num_heads //
288
- n_groups, # ratio for mapping back to original group
302
+ n_groups == 1, # if there was only one group
303
)
290
- intermediate_settings = (intermediate_size, 0, 1)
291
- head_setings = (self.num_heads, 0, 1)
304
+ intermediate_settings = (intermediate_size, 0, False)
305
+ head_setings = (self.num_heads, 0, False)
306
307
# - the weight already has a "weight_loader" attribute
308
# which set_weight_attrs will raise if we do not