shared_moe_weight_scale parameter

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

Author: jenchen13

modelopt/torch/quantization/model_calib.py

@@ -95,13 +95,20 @@ def _check_moe_calibration_complete(quantizer, parallel_state):
 
 
 @torch.no_grad()
-def max_calibrate(model: nn.Module, forward_loop: ForwardLoop | None = None, distributed_sync=True):
+def max_calibrate(
+    model: nn.Module,
+    forward_loop: ForwardLoop | None = None,
+    distributed_sync=True,
+    shared_moe_weight_scale=True,
+):
     """Calibrate the model using max.
 
     Args:
         model: Model to be calibrated.
         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.
@@ -116,7 +123,7 @@ def max_calibrate(model: nn.Module, forward_loop: ForwardLoop | None = None, dis
     # 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()
+            module.layer_sync_moe_local_experts_amax(shared_moe_weight_scale)
 
     if not distributed_sync:
         return

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):
-        """Sync input quantizer amax across local experts in a SequentialMLP.
+    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.
 
         Ensures all experts have the same input quantizer amax.This function operates
         on a single rank and does not require distributed sync.
@@ -589,33 +589,35 @@ def layer_sync_moe_local_experts_amax(self):
             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
-                    and "input_quantizer" in name
-                ):
-                    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)
-                    )
+                if isinstance(module, TensorQuantizer) and module.amax is not None:
+                    if shared_moe_weight_scale or (
+                        not shared_moe_weight_scale and "input_quantizer" 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)
+                        )
 
         # Apply synchronized amax values back to all local experts
         for expert in self.local_experts:
             for name, module in expert.named_modules():
-                if (
-                    isinstance(module, TensorQuantizer)
-                    and name in amax_dict
-                    and "input_quantizer" in name
-                ):
-                    module.amax = amax_dict[name].detach().clone()
+                if isinstance(module, TensorQuantizer) and name in amax_dict:
+                    if shared_moe_weight_scale or (
+                        not shared_moe_weight_scale and "input_quantizer" in name
+                    ):
+                        module.amax = amax_dict[name].detach().clone()
 
     def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
         """Override the default to enable singleton_local_shards.

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

@@ -737,7 +737,8 @@ 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])
-def test_layer_sync_moe_local_experts_amax(ep_size, moe_grouped_gemm):
+@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):
     """Test expert model parallel synchronization."""
     size = torch.cuda.device_count()
     if size < ep_size:
@@ -749,12 +750,15 @@ def test_layer_sync_moe_local_experts_amax(ep_size, moe_grouped_gemm):
             _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, rank, size):
+def _test_layer_sync_moe_local_experts_amax(
+    ep_size, moe_grouped_gemm, shared_moe_weight_scale, rank, size
+):
     initialize_for_megatron(
         tensor_model_parallel_size=1,
         pipeline_model_parallel_size=1,
@@ -776,7 +780,7 @@ def _test_layer_sync_moe_local_experts_amax(ep_size, moe_grouped_gemm, rank, siz
 
     # Sync amax across local experts in each layer
     for layer in model.decoder.layers:
-        layer.mlp.experts.layer_sync_moe_local_experts_amax()
+        layer.mlp.experts.layer_sync_moe_local_experts_amax(shared_moe_weight_scale)
 
     for layer in model.decoder.layers:
         fc1_amax = None
@@ -793,6 +797,22 @@ def _test_layer_sync_moe_local_experts_amax(ep_size, moe_grouped_gemm, rank, siz
             else:
                 assert torch.allclose(fc2_amax, expert.linear_fc2.input_quantizer.amax)
 
+    if shared_moe_weight_scale:
+        for layer in model.decoder.layers:
+            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
+                else:
+                    assert torch.allclose(fc1_amax, expert.linear_fc1.weight_quantizer.amax)
+                if fc2_amax is None:
+                    fc2_amax = expert.linear_fc2.weight_quantizer.amax
+                else:
+                    assert torch.allclose(fc2_amax, expert.linear_fc2.weight_quantizer.amax)
+
 
 def _test_expert_model_parallel_amax_sync(
     tp_size, ep_size, etp_size, moe_grouped_gemm, config, rank, size