move around imports

Author: dsikka

src/llmcompressor/modeling/prepare.py

@@ -1,23 +1,12 @@
 import tqdm
 from compressed_tensors.utils import replace_module
-from loguru import logger
 from transformers import PreTrainedModel
 
 from llmcompressor.modeling.deepseek_v3 import replace as replace_deepseekv3
 from llmcompressor.modeling.llama4 import replace as replace_llama4
 from llmcompressor.modeling.qwen3_moe import replace as replace_Qwen3MoE
-
-try:
-    from llmcompressor.modeling.qwen3_next_moe import replace as replace_Qwen3NextMoE
-    from llmcompressor.modeling.qwen3_vl_moe import replace as replace_Qwen3VLMoE
-except ImportError:
-    logger.warning(
-        "Qwen3-VL-MoE and Qwen3-Next support is not available. "
-        "Please ensure that you have the correct version of transformers installed."
-    )
-    replace_Qwen3VLMoE = None
-    replace_Qwen3NextMoE = None
-
+from llmcompressor.modeling.qwen3_next_moe import replace as replace_Qwen3NextMoE
+from llmcompressor.modeling.qwen3_vl_moe import replace as replace_Qwen3VLMoE
 from llmcompressor.utils.helpers import patch_attr
 
 __all__ = ["replace_modules_for_calibration"]
@@ -26,11 +15,9 @@
 replacements = {
     "DeepseekV3MoE": replace_deepseekv3,
     "Llama4TextMoe": replace_llama4,
+    "Qwen3VLMoeTextSparseMoeBlock": replace_Qwen3VLMoE,
 }
 
-if replace_Qwen3VLMoE is not None:
-    replacements["Qwen3VLMoeTextSparseMoeBlock"] = replace_Qwen3VLMoE
-
 
 def replace_modules_for_calibration(
     model: PreTrainedModel,
@@ -92,11 +79,9 @@ def update_qwen3_next_moe(model, module, stack, calibrate_all_experts):
 
 moe_context = {
     "Qwen3MoeForCausalLM": update_qwen3_moe,
+    "Qwen3NextForCausalLM": update_qwen3_next_moe,
 }
 
-if replace_Qwen3NextMoE is not None:
-    moe_context["Qwen3NextForCausalLM"] = update_qwen3_next_moe
-
 
 def moe_calibration_context(
     model: PreTrainedModel,

src/llmcompressor/modeling/qwen3_next_moe.py

@@ -15,17 +15,13 @@
 # limitations under the License.
 
 import torch
-from transformers.models import Qwen3NextConfig
-from transformers.models.qwen3_next.modeling_qwen3_next import (
-    Qwen3NextSparseMoeBlock as OriginalQwen3NextMoeSparseMoeBlock,
-)
 
 
 class Qwen3NextSparseMoeBlock(torch.nn.Module):
     def __init__(
         self,
-        config: Qwen3NextConfig,
-        original: OriginalQwen3NextMoeSparseMoeBlock,
+        config,
+        original,
         calibrate_all_experts: bool,
     ):
         super().__init__()
@@ -109,9 +105,9 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
 
 def replace(
-    config: Qwen3NextConfig,
-    module: OriginalQwen3NextMoeSparseMoeBlock,
-    calibrate_all_experts: bool,
+    config,
+    module,
+    calibrate_all_experts,
 ):
     return Qwen3NextSparseMoeBlock(
         config=config, original=module, calibrate_all_experts=calibrate_all_experts

src/llmcompressor/modeling/qwen3_vl_moe.py

@@ -1,7 +1,4 @@
 import torch
-from transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe import (
-    Qwen3VLMoeTextMLP,
-)
 
 from llmcompressor.utils.dev import skip_weights_initialize
 
@@ -17,6 +14,10 @@ def __init__(self, config, original):
 
 class SequentialQwen3VLMoeTextExperts(torch.nn.ModuleList):
     def __init__(self, config, original):
+        from transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe import (
+            Qwen3VLMoeTextMLP,
+        )
+
         self.num_experts = original.gate_up_proj.shape[0]
         with skip_weights_initialize():
             super().__init__(