[Model] Enhance error reporting for invalid tensor-parallel settings (#2566)

MasterJH5574 · web-flow · commit 873827c25ca1 · 2024-06-12T07:14:08.000-04:00
This PR enhances the error reporting for multi-GPU model compilation,
so we can provide as many error reasons as possible before loading and
running the models.
diff --git a/python/mlc_llm/model/baichuan/baichuan_model.py b/python/mlc_llm/model/baichuan/baichuan_model.py
@@ -87,6 +87,11 @@ def __post_init__(self):
 class BaichuanAttention(nn.Module):  # pylint: disable=too-many-instance-attributes
     def __init__(self, config: BaichuanConfig):
         self.hidden_size = config.hidden_size
+        if config.num_attention_heads % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split {config.num_attention_heads} attention heads "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.num_heads = config.num_attention_heads // config.tensor_parallel_shards
         self.head_dim = config.head_dim
         self.W_pack = nn.Linear(self.hidden_size, 3 * self.num_heads * self.head_dim, bias=False)
@@ -106,6 +111,11 @@ def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id:
 
 class BaichuanMLP(nn.Module):
     def __init__(self, config: BaichuanConfig):
+        if config.intermediate_size % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split MLP intermediate size {config.intermediate_size} "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards
         self.gate_up_proj = nn.Linear(
             in_features=config.hidden_size,
diff --git a/python/mlc_llm/model/baichuan/baichuan_quantization.py b/python/mlc_llm/model/baichuan/baichuan_quantization.py
@@ -19,6 +19,7 @@ def group_quant(
     model: nn.Module = BaichuanForCausalLM(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/bert/bert_model.py b/python/mlc_llm/model/bert/bert_model.py
@@ -83,6 +83,11 @@ def __post_init__(self):
 
 class BertSelfAttention(nn.Module):  # pylint: disable=too-many-instance-attributes
     def __init__(self, config: BertConfig):
+        if config.num_attention_heads % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split {config.num_attention_heads} attention heads"
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.num_heads = config.num_attention_heads // config.tensor_parallel_shards
         self.head_dim = config.head_dim
 
diff --git a/python/mlc_llm/model/bert/bert_quantization.py b/python/mlc_llm/model/bert/bert_quantization.py
@@ -19,6 +19,7 @@ def group_quant(
     model: nn.Module = BertModel(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/chatglm3/chatglm3_model.py b/python/mlc_llm/model/chatglm3/chatglm3_model.py
@@ -93,6 +93,11 @@ def __post_init__(self):
 class GLMAttention(nn.Module):  # pylint: disable=too-many-instance-attributes
     def __init__(self, config: GLMConfig):
         self.hidden_size = config.hidden_size
+        if config.num_attention_heads % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split {config.num_attention_heads} attention heads"
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.num_heads = config.num_attention_heads // config.tensor_parallel_shards
         self.multi_query_attention = config.multi_query_attention
         self.num_key_value_heads = (
@@ -125,6 +130,11 @@ def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id:
 
 class GLMMLP(nn.Module):
     def __init__(self, config: GLMConfig):
+        if config.ffn_hidden_size % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split ffn hidden size {config.ffn_hidden_size} "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.ffn_hidden_size = config.ffn_hidden_size // config.tensor_parallel_shards
 
         self.dense_h_to_4h = nn.Linear(
diff --git a/python/mlc_llm/model/chatglm3/chatglm3_quantization.py b/python/mlc_llm/model/chatglm3/chatglm3_quantization.py
@@ -19,6 +19,7 @@ def group_quant(
     model: nn.Module = ChatGLMForCausalLM(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/eagle/eagle_quantization.py b/python/mlc_llm/model/eagle/eagle_quantization.py
@@ -19,6 +19,7 @@ def group_quant(
     model: nn.Module = EagleForCasualLM(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/gemma/gemma_model.py b/python/mlc_llm/model/gemma/gemma_model.py
@@ -102,6 +102,11 @@ def lm_head_forward(self, x: nn.Tensor):
 class GemmaMLP(nn.Module):
     def __init__(self, config: GemmaConfig):
         super().__init__()
+        if config.intermediate_size % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split MLP intermediate size {config.intermediate_size} "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards
         self.gate_up_proj = nn.Linear(
             in_features=config.hidden_size,
diff --git a/python/mlc_llm/model/gemma/gemma_quantization.py b/python/mlc_llm/model/gemma/gemma_quantization.py
@@ -19,6 +19,7 @@ def group_quant(
     model: nn.Module = GemmaForCausalLM(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/gpt2/gpt2_model.py b/python/mlc_llm/model/gpt2/gpt2_model.py
@@ -84,6 +84,11 @@ def __post_init__(self):
 class GPT2Attention(nn.Module):  # pylint: disable=too-many-instance-attributes
     def __init__(self, config: GPT2Config):
         self.embed_dim = config.n_embd
+        if config.n_head % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split {config.n_head} attention heads "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.num_heads = config.n_head // config.tensor_parallel_shards
         self.head_dim = config.head_dim
         self.scale_attn_by_inverse_layer_idx = config.scale_attn_by_inverse_layer_idx
@@ -120,6 +125,11 @@ def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id:
 class GPT2MLP(nn.Module):
     def __init__(self, config: GPT2Config):
         embed_dim = config.n_embd
+        if config.n_inner % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split MLP intermediate size {config.n_inner} "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         intermediate_size = config.n_inner // config.tensor_parallel_shards
         self.c_fc = nn.Linear(embed_dim, intermediate_size)
         self.c_proj = nn.Linear(intermediate_size, embed_dim)
diff --git a/python/mlc_llm/model/gpt2/gpt2_quantization.py b/python/mlc_llm/model/gpt2/gpt2_quantization.py
@@ -19,6 +19,7 @@ def group_quant(
     model: nn.Module = GPT2LMHeadModel(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/gpt_bigcode/gpt_bigcode_quantization.py b/python/mlc_llm/model/gpt_bigcode/gpt_bigcode_quantization.py
@@ -19,6 +19,7 @@ def group_quant(
     model: nn.Module = GPTBigCodeForCausalLM(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/gpt_neox/gpt_neox_model.py b/python/mlc_llm/model/gpt_neox/gpt_neox_model.py
@@ -94,6 +94,11 @@ class GPTNeoXAttention(nn.Module):  # pylint: disable=too-many-instance-attribut
     def __init__(self, config: GPTNeoXConfig):
         self.rope_theta = config.position_embedding_base
         self.hidden_size = config.hidden_size
+        if config.num_attention_heads % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split {config.num_attention_heads} attention heads "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.num_attention_heads = config.num_attention_heads // config.tensor_parallel_shards
         self.head_dim = config.head_dim
         self.query_key_value = nn.Linear(
@@ -126,6 +131,11 @@ class GPTNeoXMLP(nn.Module):
     def __init__(self, config: GPTNeoXConfig):
         super().__init__()
         out_dtype = config.ffn_out_dtype
+        if config.intermediate_size % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split MLP intermediate size {config.intermediate_size} "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards
         self.dense_h_to_4h = nn.Linear(
             config.hidden_size,
diff --git a/python/mlc_llm/model/gpt_neox/gpt_neox_quantization.py b/python/mlc_llm/model/gpt_neox/gpt_neox_quantization.py
@@ -19,6 +19,7 @@ def group_quant(
     model: nn.Module = GPTNeoXForCausalLM(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/internlm/internlm_model.py b/python/mlc_llm/model/internlm/internlm_model.py
@@ -86,6 +86,11 @@ def __post_init__(self):
 class InternLMAttention(nn.Module):  # pylint: disable=too-many-instance-attributes
     def __init__(self, config: InternLMConfig):
         self.hidden_size = config.hidden_size
+        if config.num_attention_heads % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split {config.num_attention_heads} attention heads "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.num_heads = config.num_attention_heads // config.tensor_parallel_shards
         self.head_dim = config.head_dim
         self.max_position_embeddings = config.context_window_size
@@ -109,6 +114,11 @@ def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id:
 
 class InternLMMLP(nn.Module):
     def __init__(self, config: InternLMConfig):
+        if config.intermediate_size % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split MLP intermediate size {config.intermediate_size} "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards
 
         self.gate_up_proj = nn.Linear(
diff --git a/python/mlc_llm/model/internlm/internlm_quantization.py b/python/mlc_llm/model/internlm/internlm_quantization.py
@@ -19,6 +19,7 @@ def group_quant(
     model: nn.Module = InternLMForCausalLM(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/llama/llama_model.py b/python/mlc_llm/model/llama/llama_model.py
@@ -91,6 +91,11 @@ def __post_init__(self):
 class LlamaFFN(nn.Module):
     def __init__(self, config: LlamaConfig):
         super().__init__()
+        if config.intermediate_size % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split MLP intermediate size {config.intermediate_size} "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards
         self.gate_up_proj = nn.Linear(
             in_features=config.hidden_size,
diff --git a/python/mlc_llm/model/llama/llama_quantization.py b/python/mlc_llm/model/llama/llama_quantization.py
@@ -25,6 +25,7 @@ def group_quant(
     model: nn.Module = LlamaForCasualLM(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/llava/llava_quantization.py b/python/mlc_llm/model/llava/llava_quantization.py
@@ -18,6 +18,7 @@ def group_quant(
     model: nn.Module = LlavaForCasualLM(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/mistral/mistral_model.py b/python/mlc_llm/model/mistral/mistral_model.py
@@ -101,6 +101,11 @@ class MistralMLP(nn.Module):
 
     def __init__(self, config: MistralConfig):
         super().__init__()
+        if config.intermediate_size % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split MLP intermediate size {config.intermediate_size} "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards
         self.gate_up_proj = nn.Linear(
             in_features=config.hidden_size,
@@ -120,6 +125,11 @@ class MistralAttention(nn.Module):  # pylint: disable=too-many-instance-attribut
 
     def __init__(self, config: MistralConfig):
         self.head_dim = config.head_dim
+        if config.num_key_value_heads % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split {config.num_key_value_heads} key-value attention heads "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.num_q_heads = config.num_attention_heads // config.tensor_parallel_shards
         self.num_kv_heads = config.num_key_value_heads // config.tensor_parallel_shards
         self.qkv_proj = nn.Linear(
diff --git a/python/mlc_llm/model/mistral/mistral_quantization.py b/python/mlc_llm/model/mistral/mistral_quantization.py
@@ -19,6 +19,7 @@ def group_quant(
     model: nn.Module = MistralForCasualLM(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/mixtral/mixtral_model.py b/python/mlc_llm/model/mixtral/mixtral_model.py
@@ -39,6 +39,11 @@ def __init__(self, config: MixtralConfig):
         super().__init__()
         self.num_experts_per_tok = config.num_experts_per_tok
         self.num_local_experts = config.num_local_experts
+        if config.intermediate_size % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split MoE intermediate size {config.intermediate_size} "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards
         self.gate = nn.Linear(
             in_features=config.hidden_size,
diff --git a/python/mlc_llm/model/mixtral/mixtral_quantization.py b/python/mlc_llm/model/mixtral/mixtral_quantization.py
@@ -25,6 +25,7 @@ def group_quant(
     model: nn.Module = MixtralForCasualLM(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/orion/orion_model.py b/python/mlc_llm/model/orion/orion_model.py
@@ -91,6 +91,11 @@ def __post_init__(self):
 class OrionFFN(nn.Module):
     def __init__(self, config: OrionConfig):
         super().__init__()
+        if config.intermediate_size % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split MLP intermediate size {config.intermediate_size} "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards
         self.gate_up_proj = nn.Linear(
             in_features=config.hidden_size,
diff --git a/python/mlc_llm/model/orion/orion_quantization.py b/python/mlc_llm/model/orion/orion_quantization.py
@@ -19,6 +19,7 @@ def group_quant(
     model: nn.Module = OrionForCasualLM(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/phi/phi_model.py b/python/mlc_llm/model/phi/phi_model.py
@@ -176,6 +176,11 @@ def from_phi1(config: Phi1Config) -> "PhiConfig":
 class PhiMLP(nn.Module):
     def __init__(self, config: PhiConfig):
         super().__init__()
+        if config.n_inner % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split MLP intermediate size {config.n_inner} "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.intermediate_size = config.n_inner // config.tensor_parallel_shards
         self.fc1 = nn.Linear(config.n_embd, self.intermediate_size)
         self.fc2 = nn.Linear(self.intermediate_size, config.n_embd)
diff --git a/python/mlc_llm/model/phi/phi_quantization.py b/python/mlc_llm/model/phi/phi_quantization.py
@@ -19,6 +19,7 @@ def group_quant(
     model: nn.Module = PhiForCausalLM(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/phi3/phi3_model.py b/python/mlc_llm/model/phi3/phi3_model.py
@@ -94,6 +94,11 @@ def __post_init__(self):
 class Phi3MLP(nn.Module):
     def __init__(self, config: Phi3Config):
         super().__init__()
+        if config.intermediate_size % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split MLP intermediate size {config.intermediate_size} "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards
         self.gate_up_proj = nn.Linear(config.hidden_size, 2 * self.intermediate_size, bias=False)
         self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False)
diff --git a/python/mlc_llm/model/phi3/phi3_quantization.py b/python/mlc_llm/model/phi3/phi3_quantization.py
@@ -19,6 +19,7 @@ def group_quant(
     model: nn.Module = Phi3ForCausalLM(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/qwen/qwen_model.py b/python/mlc_llm/model/qwen/qwen_model.py
@@ -84,6 +84,11 @@ def __post_init__(self):
 class QWenAttention(nn.Module):  # pylint: disable=too-many-instance-attributes
     def __init__(self, config: QWenConfig):
         self.hidden_size = config.hidden_size
+        if config.num_attention_heads % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split {config.num_attention_heads} attention heads "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.num_heads = config.num_attention_heads // config.tensor_parallel_shards
         self.head_dim = config.head_dim
 
@@ -110,6 +115,11 @@ def forward(  # pylint: disable=too-many-locals
 
 class QWenMLP(nn.Module):
     def __init__(self, config: QWenConfig):
+        if config.intermediate_size % config.tensor_parallel_shards != 0:
+            raise ValueError(
+                f"Cannot split MLP intermediate size {config.intermediate_size} "
+                f"evenly to {config.tensor_parallel_shards} GPUs."
+            )
         self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards
         self.gate_up_proj = nn.Linear(
             in_features=config.hidden_size,
diff --git a/python/mlc_llm/model/qwen/qwen_quantization.py b/python/mlc_llm/model/qwen/qwen_quantization.py
@@ -19,6 +19,7 @@ def group_quant(
     model: nn.Module = QWenLMHeadModel(model_config)
     model.to(quantization.model_dtype)
     quant_map = QuantizeMapping({}, {})
+    quantization.tensor_parallel_shards = model_config.tensor_parallel_shards
     model = quantization.quantize_model(
         model,
         quant_map,
diff --git a/python/mlc_llm/model/qwen2/qwen2_model.py b/python/mlc_llm/model/qwen2/qwen2_model.py
diff --git a/python/mlc_llm/model/qwen2/qwen2_quantization.py b/python/mlc_llm/model/qwen2/qwen2_quantization.py
diff --git a/python/mlc_llm/model/qwen2_moe/qwen2_moe_model.py b/python/mlc_llm/model/qwen2_moe/qwen2_moe_model.py
diff --git a/python/mlc_llm/model/qwen2_moe/qwen2_moe_quantization.py b/python/mlc_llm/model/qwen2_moe/qwen2_moe_quantization.py
diff --git a/python/mlc_llm/model/rwkv5/rwkv5_model.py b/python/mlc_llm/model/rwkv5/rwkv5_model.py
diff --git a/python/mlc_llm/model/rwkv5/rwkv5_quantization.py b/python/mlc_llm/model/rwkv5/rwkv5_quantization.py
diff --git a/python/mlc_llm/model/rwkv6/rwkv6_model.py b/python/mlc_llm/model/rwkv6/rwkv6_model.py
diff --git a/python/mlc_llm/model/rwkv6/rwkv6_quantization.py b/python/mlc_llm/model/rwkv6/rwkv6_quantization.py
diff --git a/python/mlc_llm/model/stable_lm/stablelm_model.py b/python/mlc_llm/model/stable_lm/stablelm_model.py
diff --git a/python/mlc_llm/model/stable_lm/stablelm_quantization.py b/python/mlc_llm/model/stable_lm/stablelm_quantization.py
diff --git a/python/mlc_llm/quantization/group_quantization.py b/python/mlc_llm/quantization/group_quantization.py
