update configs with shared_moe_weight_scale

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

Author: jenchen13

modelopt/torch/quantization/config.py

@@ -214,7 +214,7 @@
         **_default_disabled_quantizer_cfg,
         **_mamba_moe_disabled_quantizer_cfg,
     },
-    "algorithm": "max",
+    "algorithm": {"method": "max", "shared_moe_weight_scale": False},
 }
 
 MAMBA_MOE_FP8_CONSERVATIVE_CFG = {
@@ -226,7 +226,7 @@
         "*mixer.in_proj*": {"enable": False},  # Skip mamba linear
         "*mixer.out_proj*": {"enable": False},  # Skip mamba linear
     },
-    "algorithm": "max",
+    "algorithm": {"method": "max", "shared_moe_weight_scale": False},
 }
 
 FP8_PER_CHANNEL_PER_TOKEN_CFG = {
@@ -437,7 +437,7 @@
         **_default_disabled_quantizer_cfg,
         **_mamba_moe_disabled_quantizer_cfg,
     },
-    "algorithm": "max",
+    "algorithm": {"method": "max", "shared_moe_weight_scale": False},
 }
 MAMBA_MOE_NVFP4_CONSERVATIVE_CFG = {
     "quant_cfg": {
@@ -458,7 +458,7 @@
         "*mixer.in_proj*": {"enable": False},  # Skip mamba linear
         "*mixer.out_proj*": {"enable": False},  # Skip mamba linear
     },
-    "algorithm": "max",
+    "algorithm": {"method": "max", "shared_moe_weight_scale": False},
 }
 
 
@@ -1087,6 +1087,12 @@ 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.

tests/gpu_megatron/torch/quantization/plugins/test_megatron.py

@@ -473,10 +473,7 @@ def test_homogeneous_sharded_state_dict(tmp_path, config, compress, meta_device,
 
 @pytest.mark.parametrize(
     "config",
-    [
-        NVFP4_GEMM_KV_CFG,
-        FP8_GEMM_KV_CFG,
-    ],
+    [NVFP4_GEMM_KV_CFG, FP8_GEMM_KV_CFG, mtq.MAMBA_MOE_NVFP4_CONSERVATIVE_CFG],
 )
 def test_homogeneous_sharded_state_dict_hybrid(tmp_path, config):
     """Test sharded state dict for hybrid Mamba MOE models."""