fix typing

Author: dsikka

src/llmcompressor/modeling/deepseek_v3.py

@@ -1,16 +1,15 @@
 import torch
 from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config
-from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3MoE
-
-__all__ = ["DeepseekV3MoECalibrate"]
-
+from transformers.models.deepseek_v3.modeling_deepseek_v3 import (
+    DeepseekV3MoE as OriginalDeepseekV3MoE
+)
 
 class DeepseekV3MoE(torch.nn.Module):
     """
     Patched DeepseekV3MoE which sends all tokens to all experts for calibration
     """
 
-    def __init__(self, config: DeepseekV3Config, original: DeepseekV3MoE):
+    def __init__(self, config: DeepseekV3Config, original: OriginalDeepseekV3MoE):
         super().__init__()
         self.config = config
         self.experts = original.experts
@@ -49,5 +48,5 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         return hidden_states
 
 
-def replace(config: DeepseekV3Config, module: DeepseekV3MoE):
+def replace(config: DeepseekV3Config, module: OriginalDeepseekV3MoE):
     return DeepseekV3MoECalibrate(config=config, original=module)

src/llmcompressor/modeling/prepare.py

@@ -32,6 +32,7 @@ def update_qwen3_moe(model, stack):
     for module in model.modules():
         cls_name = module.__class__.__name__
         if cls_name == "Qwen3MoeDecoderLayer":
+            # Optionally update the model.config to pass in other arguments
             stack.enter_context(
                 patch_attr(module, "mlp", replace_Qwen3MoE(model.config, module.mlp))
             )