[Mypy] Fix mypy typing for Zamba2 models (#58255)
Signed-off-by: Ashraf Bhuiyan <mbhuiyan@redhat.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Ashraf Bhuiyan committed
Sep 29, 2026 at 14:13 UTC
9a190f142c8acccffd0b19d2393e1eb76f775aea
2 files changed
+15
-11
tools/pre_commit/mypy.py
-1
@@ -113,7 +113,6 @@ SEPARATE_GROUPS = [
113
114
EXCLUDE = [
115
r"vllm/model_executor/models/[kK]",
116
- r"vllm/model_executor/models/[zZ]",
116
]
117
118
vllm/model_executor/models/zamba2.py
+15
-10
@@ -86,16 +86,21 @@ class Zamba2LoRA(nn.Module):
86
)
87
88
if isinstance(output_dim, list):
89
- B_class = MergedColumnParallelLinear
89
+ self.B: ColumnParallelLinear = MergedColumnParallelLinear(
90
+ rank,
91
+ output_dim,
92
+ bias=False,
93
+ quant_config=quant_config,
94
+ prefix=f"{prefix}.B",
95
+ )
96
else:
91
- B_class = ColumnParallelLinear
92
- self.B = B_class(
93
- rank,
94
- output_dim,
95
- bias=False,
96
- quant_config=quant_config,
97
- prefix=f"{prefix}.B",
98
- )
97
+ self.B = ColumnParallelLinear(
98
+ rank,
99
+ output_dim,
100
+ bias=False,
101
+ quant_config=quant_config,
102
+ prefix=f"{prefix}.B",
103
+ )
104
105
def forward(
106
self,
@@ -304,7 +309,7 @@ class Zamba2MLP(nn.Module):
309
self,
310
config: Zamba2Config,
311
bare_block_idx: int,
307
- num_hybrid_layers: dict[int, int],
312
+ num_hybrid_layers: int,
313
quant_config: QuantizationConfig | None = None,
314
prefix: str = "",
315
) -> None: