test(moe): add GLM MegaMoE validation coverage
Add parser coverage for --moe-a2a-backend=megamoe and local unit coverage for MegaMoE gating, GLM fast-path routing, fallback routing, and sidecar weight preservation. Use monkeypatch-style tests so the forward and weight-layout checks do not require loading GLM 5.2 or importing a real deep_gemm runtime. Constraint: tests must not stage or depend on local skill/doc artifacts. Feature-flag: --moe-a2a-backend=megamoe. Conflict-hotspots: test/registered/unit/server_args/test_server_args.py, test/registered/unit/moe/test_glm_megamoe.py. Scope-risk: runtime GPU e2e still needs target Blackwell/DeepGEMM environment. Tested: PYTHONPYCACHEPREFIX=/private/tmp/sglang_pycache python3 -m py_compile test/registered/unit/server_args/test_server_args.py test/registered/unit/moe/test_glm_megamoe.py. Tested: git diff --check. Not-tested: PYTHONPATH=python python3 -m pytest test/registered/unit/moe/test_glm_megamoe.py -q; local Python has no pytest. Not-tested: PYTHONPATH=python python3 -m unittest test.registered.unit.moe.test_glm_megamoe -v; local Python has no torch. Not-tested: GLM 5.2 MegaMoE GPU e2e; local environment lacks target runtime and hardware.
This commit is contained in:
LuminolT
140
test/registered/unit/moe/test_glm_megamoe.py
Normal file
140
test/registered/unit/moe/test_glm_megamoe.py
Normal file
@@ -0,0 +1,140 @@
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.moe import mega_moe
|
||||
from sglang.srt.models import glm4_moe
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
|
||||
|
||||
|
||||
class _MegaBackend:
|
||||
def is_megamoe(self):
|
||||
return True
|
||||
|
||||
|
||||
class _NonDeepEPBackend:
|
||||
def is_deepep(self):
|
||||
return False
|
||||
|
||||
|
||||
class _Param:
|
||||
def __init__(self, data):
|
||||
self.data = data
|
||||
|
||||
|
||||
class TestGLMMegaMoE(unittest.TestCase):
|
||||
def test_should_use_mega_moe_respects_env_and_token_cap(self):
|
||||
hidden_states = torch.empty((2, 32))
|
||||
moe = SimpleNamespace(
|
||||
experts=SimpleNamespace(_mega_moe_weights_built=True),
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
mega_moe, "get_moe_a2a_backend", return_value=_MegaBackend()
|
||||
), patch.object(
|
||||
mega_moe, "get_dp_global_num_tokens", return_value=None
|
||||
), patch.object(
|
||||
mega_moe, "get_is_capture_mode", return_value=False
|
||||
), envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.override(False):
|
||||
self.assertFalse(mega_moe.should_use_mega_moe(moe, hidden_states))
|
||||
|
||||
with patch.object(
|
||||
mega_moe, "get_moe_a2a_backend", return_value=_MegaBackend()
|
||||
), patch.object(
|
||||
mega_moe, "get_dp_global_num_tokens", return_value=None
|
||||
), patch.object(
|
||||
mega_moe, "get_is_capture_mode", return_value=False
|
||||
), envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.override(
|
||||
True
|
||||
), envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.override(1):
|
||||
self.assertFalse(mega_moe.should_use_mega_moe(moe, hidden_states))
|
||||
|
||||
with patch.object(
|
||||
mega_moe, "get_moe_a2a_backend", return_value=_MegaBackend()
|
||||
), patch.object(
|
||||
mega_moe, "get_dp_global_num_tokens", return_value=None
|
||||
), patch.object(
|
||||
mega_moe, "get_is_capture_mode", return_value=False
|
||||
), envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_USE_FP4_ACTS.override(
|
||||
True
|
||||
), envs.SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK.override(2):
|
||||
self.assertTrue(mega_moe.should_use_mega_moe(moe, hidden_states))
|
||||
|
||||
def test_glm_forward_uses_megamoe_fast_path(self):
|
||||
block = object.__new__(glm4_moe.Glm4MoeSparseMoeBlock)
|
||||
hidden_states = torch.empty((1, 32))
|
||||
forward_batch = object()
|
||||
expected = torch.empty_like(hidden_states)
|
||||
|
||||
with patch.object(
|
||||
mega_moe, "should_use_mega_moe", return_value=True
|
||||
) as should, patch.object(
|
||||
mega_moe, "forward_mega_moe", return_value=expected
|
||||
) as forward:
|
||||
output = block.forward(hidden_states, forward_batch=forward_batch)
|
||||
|
||||
self.assertIs(output, expected)
|
||||
should.assert_called_once_with(block, hidden_states)
|
||||
forward.assert_called_once_with(block, hidden_states, forward_batch)
|
||||
|
||||
def test_glm_forward_falls_back_to_normal_path(self):
|
||||
block = object.__new__(glm4_moe.Glm4MoeSparseMoeBlock)
|
||||
object.__setattr__(block, "alt_stream", None)
|
||||
object.__setattr__(block, "num_fused_shared_experts", 0)
|
||||
object.__setattr__(block, "forward_normal", Mock(return_value="normal-output"))
|
||||
hidden_states = torch.empty((1, 32))
|
||||
|
||||
with patch.object(
|
||||
mega_moe, "should_use_mega_moe", return_value=False
|
||||
), patch.object(
|
||||
glm4_moe, "get_moe_a2a_backend", return_value=_NonDeepEPBackend()
|
||||
):
|
||||
output = block.forward(hidden_states)
|
||||
|
||||
self.assertEqual(output, "normal-output")
|
||||
block.forward_normal.assert_called_once_with(hidden_states, False, False)
|
||||
|
||||
def test_build_mega_moe_weights_preserves_runner_layout(self):
|
||||
def fake_transform_sf(sf, *, mn, k, recipe, num_groups, disable_ue8m0_cast):
|
||||
del sf, k, recipe, disable_ue8m0_cast
|
||||
return torch.zeros((num_groups, mn, 1), dtype=torch.int32)
|
||||
|
||||
fake_deep_gemm = types.SimpleNamespace(
|
||||
transform_sf_into_required_layout=fake_transform_sf,
|
||||
)
|
||||
experts = SimpleNamespace(
|
||||
w13_weight=_Param(torch.arange(2 * 256 * 16).reshape(2, 256, 16)),
|
||||
w2_weight=_Param(torch.arange(2 * 128 * 16).reshape(2, 128, 16)),
|
||||
w13_weight_scale=_Param(torch.ones((2, 256, 1), dtype=torch.float32)),
|
||||
w2_weight_scale=_Param(torch.ones((2, 128, 1), dtype=torch.float32)),
|
||||
)
|
||||
original_w13 = experts.w13_weight.data.clone()
|
||||
|
||||
with patch.dict(sys.modules, {"deep_gemm": fake_deep_gemm}):
|
||||
mega_moe.build_mega_moe_experts_weights(
|
||||
experts, preserve_runner_layout=True
|
||||
)
|
||||
|
||||
self.assertTrue(experts._mega_moe_weights_built)
|
||||
self.assertTrue(experts._mega_moe_preserve_runner_layout)
|
||||
self.assertTrue(torch.equal(experts.w13_weight.data, original_w13))
|
||||
self.assertNotEqual(
|
||||
experts.mega_l1_weights[0].data_ptr(),
|
||||
experts.w13_weight.data.data_ptr(),
|
||||
)
|
||||
|
||||
experts.w13_weight.data.zero_()
|
||||
self.assertFalse(
|
||||
torch.equal(experts.mega_l1_weights[0], experts.w13_weight.data)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,4 +1,5 @@
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
@@ -195,6 +196,46 @@ def test_nsa_index_skip_topk_offset_zero_means_unset():
|
||||
assert override_args == {"index_topk_freq": 4}
|
||||
|
||||
|
||||
def test_megamoe_a2a_backend_parser_sets_ep_to_tp():
|
||||
import argparse
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
|
||||
field = envs.SGLANG_OPT_FIX_MEGA_MOE_MEMORY
|
||||
backup_present = field.name in os.environ
|
||||
backup_value = os.environ.get(field.name)
|
||||
backup_set_to_none = field._set_to_none
|
||||
os.environ.pop(field.name, None)
|
||||
field._set_to_none = False
|
||||
field._cache_valid = False
|
||||
|
||||
try:
|
||||
parser = argparse.ArgumentParser()
|
||||
ServerArgs.add_cli_args(parser)
|
||||
raw_args = parser.parse_args(
|
||||
[
|
||||
"--model-path",
|
||||
"dummy",
|
||||
"--tp-size",
|
||||
"4",
|
||||
"--moe-a2a-backend",
|
||||
"megamoe",
|
||||
]
|
||||
)
|
||||
args = ServerArgs.from_cli_args(raw_args)
|
||||
assert args.moe_a2a_backend == "megamoe"
|
||||
assert args.tp_size == 4
|
||||
assert args.ep_size == 4
|
||||
assert field.get() is True
|
||||
finally:
|
||||
if backup_present:
|
||||
os.environ[field.name] = backup_value
|
||||
else:
|
||||
os.environ.pop(field.name, None)
|
||||
field._set_to_none = backup_set_to_none
|
||||
field._cache_valid = False
|
||||
|
||||
|
||||
class TestLoadBalanceMethod(unittest.TestCase):
|
||||
def test_non_pd_defaults_to_round_robin(self):
|
||||
server_args = ServerArgs(model_path="dummy", disaggregation_mode="null")
|
||||
|
||||
Reference in New Issue
Block a user