[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>

tools/pre_commit/mypy.py

@@ -113,7 +113,6 @@ SEPARATE_GROUPS = [
 
 EXCLUDE = [
     r"vllm/model_executor/models/[kK]",
-    r"vllm/model_executor/models/[zZ]",
 ]
 
 

vllm/model_executor/models/zamba2.py

@@ -86,16 +86,21 @@ class Zamba2LoRA(nn.Module):
         )
 
         if isinstance(output_dim, list):
-            B_class = MergedColumnParallelLinear
+            self.B: ColumnParallelLinear = MergedColumnParallelLinear(
+                rank,
+                output_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.B",
+            )
         else:
-            B_class = ColumnParallelLinear
-        self.B = B_class(
-            rank,
-            output_dim,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.B",
-        )
+            self.B = ColumnParallelLinear(
+                rank,
+                output_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.B",
+            )
 
     def forward(
         self,
@@ -304,7 +309,7 @@ class Zamba2MLP(nn.Module):
         self,
         config: Zamba2Config,
         bare_block_idx: int,
-        num_hybrid_layers: dict[int, int],
+        num_hybrid_layers: int,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> None: