fix bug

Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>

Author: jenchen13

modelopt/torch/quantization/config.py

@@ -214,6 +214,7 @@
         **_default_disabled_quantizer_cfg,
         **_mamba_moe_disabled_quantizer_cfg,
     },
+    "algorithm": "max",
 }
 
 MAMBA_MOE_FP8_CONSERVATIVE_CFG = {
@@ -225,6 +226,7 @@
         "*mixer.in_proj*": {"enable": False},  # Skip mamba linear
         "*mixer.out_proj*": {"enable": False},  # Skip mamba linear
     },
+    "algorithm": "max",
 }
 
 FP8_PER_CHANNEL_PER_TOKEN_CFG = {
@@ -435,6 +437,7 @@
         **_default_disabled_quantizer_cfg,
         **_mamba_moe_disabled_quantizer_cfg,
     },
+    "algorithm": "max",
 }
 MAMBA_MOE_NVFP4_CONSERVATIVE_CFG = {
     "quant_cfg": {
@@ -455,6 +458,7 @@
         "*mixer.in_proj*": {"enable": False},  # Skip mamba linear
         "*mixer.out_proj*": {"enable": False},  # Skip mamba linear
     },
+    "algorithm": "max",
 }
 
 
@@ -1083,12 +1087,6 @@ class MaxCalibConfig(QuantizeAlgorithmConfig):
         description="If True, the amax will be synced across the distributed processes.",
     )
 
-    shared_moe_weight_scale: bool | None = ModeloptField(
-        default=True,
-        title="Whether to share the weight scale across local experts.",
-        description="If True, the weight scale will be shared across local experts.",
-    )
-
 
 class MseCalibConfig(QuantizeAlgorithmConfig):
     """Configuration for per-tensor MSE calibration.

modelopt/torch/quantization/model_calib.py

@@ -99,7 +99,6 @@ def max_calibrate(
     model: nn.Module,
     forward_loop: ForwardLoop | None = None,
     distributed_sync=True,
-    shared_moe_weight_scale=True,
 ):
     """Calibrate the model using max.
 
@@ -108,7 +107,6 @@ def max_calibrate(
         forward_loop: A callable which takes the model as argument and
             forwards calibration data through the model.
         distributed_sync: Whether to sync amax across distributed processes.
-        shared_moe_weight_scale: Whether to share the weight scale across local experts.
 
     See :class:`MaxCalibConfig <modelopt.torch.quantization.config.MaxCalibConfig>` for
     details on the remaining arguments.
@@ -123,7 +121,7 @@ def max_calibrate(
     # Sync amax across local experts within each rank (for SequentialMLP)
     for name, module in model.named_modules():
         if hasattr(module, "layer_sync_moe_local_experts_amax"):
-            module.layer_sync_moe_local_experts_amax(shared_moe_weight_scale)
+            module.layer_sync_moe_local_experts_amax()
 
     if not distributed_sync:
         return

modelopt/torch/quantization/plugins/megatron.py

@@ -594,7 +594,11 @@ def layer_sync_moe_local_experts_amax(self):
         amax_dict = {}
         for expert in self.local_experts:
             for name, module in expert.named_modules():
-                if isinstance(module, TensorQuantizer) and module.amax is not None:
+                if (
+                    isinstance(module, TensorQuantizer)
+                    and module.amax is not None
+                    and "input_quantizer" in name
+                ):
                     stored_amax = amax_dict.get(name)
                     amax_tensor = module.amax.detach().clone()
                     amax_dict[name] = (