@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / aaaec52ad9

Fork 0

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, 18:30 UTC

[Bugfix][Model] Mixtral: use unused head_dim config argument (#14961)

Signed-off-by: Quentin Torroba <quentin.torroba@mistral.ai>

Quentin committed Mar 17, 2025 at 15:44 UTC aaaec52ad9073ee5b80bced5d848fbff27ca64c0

2 files changed +10 -2

vllm/model_executor/models/mixtral.py

+5 -1

@@ -111,6 +111,7 @@ class MixtralAttention(nn.Module):

111

112 def __init__(

113 self,

114 + config: MixtralConfig,

115 hidden_size: int,

116 num_heads: int,

117 num_kv_heads: int,

@@ -136,7 +137,9 @@ class MixtralAttention(nn.Module):

137 # the KV heads across multiple tensor parallel GPUs.

138 assert tp_size % self.total_num_kv_heads == 0

139 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

139 - self.head_dim = hidden_size // self.total_num_heads

140 + # MixtralConfig has an optional head_dim argument

141 + self.head_dim = getattr(config, "head_dim",

142 + self.hidden_size // self.total_num_heads)

143 self.q_size = self.num_heads * self.head_dim

144 self.kv_size = self.num_kv_heads * self.head_dim

145 self.scaling = self.head_dim**-0.5

@@ -200,6 +203,7 @@ class MixtralDecoderLayer(nn.Module):

203 # Requires transformers > 4.32.0

204 rope_theta = getattr(config, "rope_theta", 10000)

205 self.self_attn = MixtralAttention(

206 + config=config,

207 hidden_size=self.hidden_size,

208 num_heads=config.num_attention_heads,

209 max_position=config.max_position_embeddings,

vllm/model_executor/models/mixtral_quant.py

+5 -1

@@ -165,6 +165,7 @@ class MixtralAttention(nn.Module):

165

166 def __init__(

167 self,

168 + config: MixtralConfig,

169 hidden_size: int,

170 num_heads: int,

171 num_kv_heads: int,

@@ -190,7 +191,9 @@ class MixtralAttention(nn.Module):

191 # the KV heads across multiple tensor parallel GPUs.

192 assert tp_size % self.total_num_kv_heads == 0

193 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

193 - self.head_dim = hidden_size // self.total_num_heads

194 + # MixtralConfig has an optional head_dim argument

195 + self.head_dim = getattr(config, "head_dim",

196 + self.hidden_size // self.total_num_heads)

197 self.q_size = self.num_heads * self.head_dim

198 self.kv_size = self.num_kv_heads * self.head_dim

199 self.scaling = self.head_dim**-0.5

@@ -252,6 +255,7 @@ class MixtralDecoderLayer(nn.Module):

255 # Requires transformers > 4.32.0

256 rope_theta = getattr(config, "rope_theta", 10000)

257 self.self_attn = MixtralAttention(

258 + config=config,

259 hidden_size=self.hidden_size,

260 num_heads=config.num_attention_heads,

261 max_position=config.max_position_embeddings,