fix test

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

Author: jenchen13

modelopt/torch/quantization/plugins/megatron.py

@@ -594,10 +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():
+                print(name, module)
                 if (
                     isinstance(module, TensorQuantizer)
                     and module.amax is not None
-                    and name == "input_quantizer"
+                    and "input_quantizer" in name
                 ):
                     stored_amax = amax_dict.get(name)
                     amax_tensor = module.amax.detach().clone()
@@ -613,7 +614,7 @@ def layer_sync_moe_local_experts_amax(self):
                 if (
                     isinstance(module, TensorQuantizer)
                     and name in amax_dict
-                    and name == "input_quantizer"
+                    and "input_quantizer" in name
                 ):
                     module.amax = amax_dict[name].detach().clone()
 

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

@@ -735,37 +735,63 @@ def test_te_grouped_vs_sequential_quantize(need_4_gpus):
     )
 
 
-def test_layer_sync_moe_local_experts_amax(moe_grouped_gemm):
+@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):
+    """Test expert model parallel synchronization."""
+    size = torch.cuda.device_count()
+    if size < ep_size:
+        pytest.skip(f"Requires at least {ep_size} GPUs for expert model parallel test")
+
+    spawn_multiprocess_job(
+        size=size,
+        job=partial(
+            _test_layer_sync_moe_local_experts_amax,
+            ep_size,
+            moe_grouped_gemm,
+        ),
+        backend="nccl",
+    )
+
+
+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,
-        expert_model_parallel_size=2,
+        expert_model_parallel_size=ep_size,
         expert_tensor_parallel_size=1,
         seed=SEED,
     )
     model = _gpt_model_provider(
         tp_size=1,
-        ep_size=2,
+        ep_size=ep_size,
         etp_size=1,
         hidden_size=256,
         moe_grouped_gemm=moe_grouped_gemm,
         use_te=moe_grouped_gemm,
         num_moe_experts=8,
         transformer_impl="modelopt",
     )
-    # model = mtq.quantize(model, mtq.FP8_DEFAULT_CFG, get_forward(model))
-    forward = get_forward(model)
-    forward()
-    print(model)
-
-    model.layer_sync_moe_local_experts_amax()
-    prev_amax = None
-    for expert in model.local_experts:
-        assert expert.input_quantizer.amax is not None
-        if prev_amax is None:
-            prev_amax = expert.input_quantizer.amax
-        else:
-            assert torch.allclose(prev_amax, expert.input_quantizer.amax)
+    model = mtq.quantize(model, mtq.FP8_DEFAULT_CFG, get_forward(model))
+
+    # Sync amax across local experts in each layer
+    for layer in model.decoder.layers:
+        layer.mlp.experts.layer_sync_moe_local_experts_amax()
+
+    for layer in model.decoder.layers:
+        fc1_amax = None
+        fc2_amax = None
+        for expert in layer.mlp.experts.local_experts:
+            assert expert.linear_fc1.input_quantizer.amax is not None
+            assert expert.linear_fc2.input_quantizer.amax is not None
+            if fc1_amax is None:
+                fc1_amax = expert.linear_fc1.input_quantizer.amax
+            else:
+                assert torch.allclose(fc1_amax, expert.linear_fc1.input_quantizer.amax)
+            if fc2_amax is None:
+                fc2_amax = expert.linear_fc2.input_quantizer.amax
+            else:
+                assert torch.allclose(fc2_amax, expert.linear_fc2.input_quantizer.amax)
 
 
 def _test_expert_model_parallel_amax_sync(
@@ -815,6 +841,7 @@ def _test_expert_model_parallel_amax_sync(
 
     # quantize the model
     model = mtq.quantize(model, config, forward)
+
     # Check initial sync status
     initial_sync, quantizer_type, rank_values = compare_amax_sync_across_expert_parallel(model)
     assert initial_sync, (