remove shared_moe flag

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

Author: jenchen13

modelopt/torch/quantization/config.py

@@ -214,7 +214,6 @@
         **_default_disabled_quantizer_cfg,
         **_mamba_moe_disabled_quantizer_cfg,
     },
-    "algorithm": {"method": "max", "shared_moe_weight_scale": False},
 }
 
 MAMBA_MOE_FP8_CONSERVATIVE_CFG = {
@@ -226,7 +225,6 @@
         "*mixer.in_proj*": {"enable": False},  # Skip mamba linear
         "*mixer.out_proj*": {"enable": False},  # Skip mamba linear
     },
-    "algorithm": {"method": "max", "shared_moe_weight_scale": False},
 }
 
 FP8_PER_CHANNEL_PER_TOKEN_CFG = {
@@ -437,7 +435,6 @@
         **_default_disabled_quantizer_cfg,
         **_mamba_moe_disabled_quantizer_cfg,
     },
-    "algorithm": {"method": "max", "shared_moe_weight_scale": False},
 }
 MAMBA_MOE_NVFP4_CONSERVATIVE_CFG = {
     "quant_cfg": {
@@ -458,7 +455,6 @@
         "*mixer.in_proj*": {"enable": False},  # Skip mamba linear
         "*mixer.out_proj*": {"enable": False},  # Skip mamba linear
     },
-    "algorithm": {"method": "max", "shared_moe_weight_scale": False},
 }
 
 

modelopt/torch/quantization/plugins/megatron.py

@@ -574,8 +574,8 @@ def _setup(self):
             expert.linear_fc1.parallel_state = self.parallel_state
             expert.linear_fc2.parallel_state = self.parallel_state
 
-    def layer_sync_moe_local_experts_amax(self, shared_moe_weight_scale=True):
-        """Sync input quantizer amax across local experts in a SequentialMLP, and optionally weight scale.
+    def layer_sync_moe_local_experts_amax(self):
+        """Sync input quantizer amax across local experts in a SequentialMLP.
 
         Ensures all experts have the same input quantizer amax.This function operates
         on a single rank and does not require distributed sync.
@@ -589,24 +589,19 @@ def layer_sync_moe_local_experts_amax(self, shared_moe_weight_scale=True):
             We need to guarantee that all experts must have amax. Otherwise, there will be deadlock
             when synchronizing over EP since some ranks may have amax None and not calling the collective
             communication.
-
-        Args:
-            shared_moe_weight_scale: Whether to share the weight scale across local experts.
         """
         # Collect amax from all local experts
         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 shared_moe_weight_scale or ("weight_quantizer" not in name):
-                        # Sync both quantizers or only sync input quantizer
-                        stored_amax = amax_dict.get(name)
-                        amax_tensor = module.amax.detach().clone()
-                        amax_dict[name] = (
-                            amax_tensor
-                            if stored_amax is None
-                            else torch.maximum(stored_amax, amax_tensor)
-                        )
+                    stored_amax = amax_dict.get(name)
+                    amax_tensor = module.amax.detach().clone()
+                    amax_dict[name] = (
+                        amax_tensor
+                        if stored_amax is None
+                        else torch.maximum(stored_amax, amax_tensor)
+                    )
 
         # Apply synchronized amax values back to all local experts
         for expert in self.local_experts:

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

@@ -734,8 +734,7 @@ def test_te_grouped_vs_sequential_quantize(need_4_gpus):
 
 @pytest.mark.parametrize("ep_size", [1, 2])
 @pytest.mark.parametrize("moe_grouped_gemm", [True, False])
-@pytest.mark.parametrize("shared_moe_weight_scale", [True, False])
-def test_layer_sync_moe_local_experts_amax(ep_size, moe_grouped_gemm, shared_moe_weight_scale):
+def test_layer_sync_moe_local_experts_amax(ep_size, moe_grouped_gemm):
     """Test expert model parallel synchronization."""
     size = torch.cuda.device_count()
     if size < ep_size:
@@ -747,15 +746,12 @@ def test_layer_sync_moe_local_experts_amax(ep_size, moe_grouped_gemm, shared_moe
             _test_layer_sync_moe_local_experts_amax,
             ep_size,
             moe_grouped_gemm,
-            shared_moe_weight_scale,
         ),
         backend="nccl",
     )
 
 
-def _test_layer_sync_moe_local_experts_amax(
-    ep_size, moe_grouped_gemm, shared_moe_weight_scale, rank, size
-):
+def _test_layer_sync_moe_local_experts_amax(ep_size, moe_grouped_gemm, rank, size):
     initialize_for_megatron(
         tensor_model_parallel_size=1,
         pipeline_model_parallel_size=1,
@@ -774,16 +770,14 @@ def _test_layer_sync_moe_local_experts_amax(
         transformer_impl="modelopt",
     )
     quant_cfg = mtq.FP8_DEFAULT_CFG
-    if not shared_moe_weight_scale:
-        quant_cfg = copy.deepcopy(quant_cfg)
-        quant_cfg["algorithm"] = {"method": "max", "shared_moe_weight_scale": False}
     model = mtq.quantize(model, quant_cfg, get_forward(model))
 
     # does layer_sync_moe_local_experts_amax happens in mtq.quantize if EP=1?
     for layer in model.decoder.layers:
-        layer.mlp.experts.layer_sync_moe_local_experts_amax(shared_moe_weight_scale)
+        layer.mlp.experts.layer_sync_moe_local_experts_amax()
 
     for layer in model.decoder.layers:
+        # Check input quantizer amax is synced across local experts
         fc1_amax = None
         fc2_amax = None
         for expert in layer.mlp.experts.local_experts:
@@ -798,25 +792,20 @@ def _test_layer_sync_moe_local_experts_amax(
             else:
                 assert torch.allclose(fc2_amax, expert.linear_fc2.input_quantizer.amax)
 
-    for layer in model.decoder.layers:
+        # Check weight quantizer amax is different across local experts
         fc1_amax = None
         fc2_amax = None
         for expert in layer.mlp.experts.local_experts:
             assert expert.linear_fc1.weight_quantizer.amax is not None
             assert expert.linear_fc2.weight_quantizer.amax is not None
             if fc1_amax is None:
                 fc1_amax = expert.linear_fc1.weight_quantizer.amax
-            elif shared_moe_weight_scale:
-                assert torch.allclose(fc1_amax, expert.linear_fc1.weight_quantizer.amax)
             else:
                 assert not torch.allclose(fc1_amax, expert.linear_fc1.weight_quantizer.amax)
-                fc1_amax = expert.linear_fc1.weight_quantizer.amax  # update most recent amax
-
             if fc2_amax is None:
                 fc2_amax = expert.linear_fc2.weight_quantizer.amax
-            elif shared_moe_weight_scale:
-                assert torch.allclose(fc2_amax, expert.linear_fc2.weight_quantizer.amax)
-            # FC2 amaxes are the same since the input to the layer is all the same
+            else:
+                assert not torch.allclose(fc2_amax, expert.linear_fc2.weight_quantizer.amax)
 
 
 def _test_expert_model_parallel_amax_sync(