From 5aa699b5909008a0d7522e0b4dec68001e0ba38a Mon Sep 17 00:00:00 2001 From: brian Date: Fri, 10 Jul 2026 17:45:15 +0800 Subject: [PATCH] fix(weight-update): preserve grouped moe expert axis --- .github/workflows/pr-test.yml | 4 + .github/workflows/pr-test.yml.j2 | 1 + .../megatron_utils/update_weight/common.py | 15 +- tests/test_grouped_moe_glu_rechunk.py | 163 ++++++++++++++++++ 4 files changed, 177 insertions(+), 6 deletions(-) create mode 100644 tests/test_grouped_moe_glu_rechunk.py diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 761496b949..e18a0dcb18 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -641,6 +641,10 @@ jobs: "num_gpus": 0, "test_file": "test_empty_colocated_weight_bucket.py" }, + { + "num_gpus": 0, + "test_file": "test_grouped_moe_glu_rechunk.py" + }, { "num_gpus": 0, "test_file": "utils/test_hf_checkpoint_saver.py" diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index 415586cb84..1de6e3667c 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -85,6 +85,7 @@ {'test_file': 'test_placement_group.py', 'num_gpus': 0}, {'test_file': 'test_external_sglang_engines.py', 'num_gpus': 0}, {'test_file': 'test_empty_colocated_weight_bucket.py', 'num_gpus': 0}, + {'test_file': 'test_grouped_moe_glu_rechunk.py', 'num_gpus': 0}, {'test_file': 'utils/test_hf_checkpoint_saver.py', 'num_gpus': 0}, {'test_file': 'plugin_contracts/test_plugin_rollout_contracts.py', 'num_gpus': 0}, {'test_file': 'plugin_contracts/test_plugin_runtime_hook_contracts.py', 'num_gpus': 0}, diff --git a/slime/backends/megatron_utils/update_weight/common.py b/slime/backends/megatron_utils/update_weight/common.py index af8715b082..91917c5efc 100644 --- a/slime/backends/megatron_utils/update_weight/common.py +++ b/slime/backends/megatron_utils/update_weight/common.py @@ -12,6 +12,11 @@ from slime.utils.types import ParamInfo +def _requires_glu_rechunk(name: str, tensor: torch.Tensor) -> bool: + """Return whether a tensor uses the dense or per-expert GLU partition layout.""" + return ("linear_fc1.weight" in name and tensor.dim() == 2) or ("linear_fc1.bias" in name and tensor.dim() == 1) + + def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor: """ All-gather TP-sharded param to full tensor. expert_bias→param, non-TP/duplicated→param.data. @@ -37,9 +42,9 @@ def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor: assert param.partition_stride == 1 or ( param.partition_stride == 2 and "linear_fc1" in name ), "partition_stride != 1 is not supported" - # TODO: here we did an extra copy during concat, maybe merge this with convert_to_hf is better? - # TODO: check only GLU is used. - if "linear_fc1.weight" in name or "linear_fc1.bias" in name: + # GLU rechunking applies to 2D dense and per-expert weights. Grouped MoE + # tensors are 3D [experts, 2 * ffn, hidden], where dim 0 is the expert axis. + if _requires_glu_rechunk(name, param): param_partitions = [p.chunk(2, dim=0) for p in param_partitions] param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions] # this is bug in megatron's grouped moe. @@ -99,9 +104,7 @@ def all_gather_params_async( else: # Process the gathered partitions (same logic as original all_gather_param) assert partition_dim is not None, "partition_stride != 1 is not supported" - # TODO: here we did an extra copy during concat, maybe merge this with convert_to_hf is better? - # TODO: check only GLU is used. - if "linear_fc1.weight" in info.name or "linear_fc1.bias" in info.name: + if _requires_glu_rechunk(info.name, param_partitions[0]): param_partitions = [p.chunk(2, dim=0) for p in param_partitions] param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions] # this is bug in megatron's grouped moe. diff --git a/tests/test_grouped_moe_glu_rechunk.py b/tests/test_grouped_moe_glu_rechunk.py new file mode 100644 index 0000000000..8e55997773 --- /dev/null +++ b/tests/test_grouped_moe_glu_rechunk.py @@ -0,0 +1,163 @@ +import importlib.util +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def _package(name, path): + module = types.ModuleType(name) + module.__path__ = [str(path)] + return module + + +def _load_common(monkeypatch): + slime_pkg = _package("slime", REPO_ROOT / "slime") + backends_pkg = _package("slime.backends", REPO_ROOT / "slime" / "backends") + megatron_utils_pkg = _package("slime.backends.megatron_utils", REPO_ROOT / "slime" / "backends" / "megatron_utils") + update_weight_pkg = _package( + "slime.backends.megatron_utils.update_weight", + REPO_ROOT / "slime" / "backends" / "megatron_utils" / "update_weight", + ) + slime_utils_pkg = _package("slime.utils", REPO_ROOT / "slime" / "utils") + + misc_utils_mod = types.ModuleType("slime.backends.megatron_utils.misc_utils") + misc_utils_mod.strip_param_name_prefix = lambda name: name.removeprefix("module.") + types_mod = types.ModuleType("slime.utils.types") + types_mod.ParamInfo = object + + megatron_mod = types.ModuleType("megatron") + core_mod = types.ModuleType("megatron.core") + mpu_mod = types.ModuleType("megatron.core.mpu") + transformer_mod = types.ModuleType("megatron.core.transformer") + transformer_layer_mod = types.ModuleType("megatron.core.transformer.transformer_layer") + transformer_layer_mod.get_transformer_layer_offset = lambda *_args, **_kwargs: 0 + core_mod.mpu = mpu_mod + + modules = { + "slime": slime_pkg, + "slime.backends": backends_pkg, + "slime.backends.megatron_utils": megatron_utils_pkg, + "slime.backends.megatron_utils.update_weight": update_weight_pkg, + "slime.backends.megatron_utils.misc_utils": misc_utils_mod, + "slime.utils": slime_utils_pkg, + "slime.utils.types": types_mod, + "megatron": megatron_mod, + "megatron.core": core_mod, + "megatron.core.mpu": mpu_mod, + "megatron.core.transformer": transformer_mod, + "megatron.core.transformer.transformer_layer": transformer_layer_mod, + } + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + + module_name = "test_grouped_moe_glu_rechunk_common" + module_path = REPO_ROOT / "slime" / "backends" / "megatron_utils" / "update_weight" / "common.py" + sys.modules.pop(module_name, None) + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def common(monkeypatch): + return _load_common(monkeypatch) + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("name", "shape", "expected"), + [ + ("decoder.layers.0.mlp.linear_fc1.weight", (1024, 2048), True), + ("decoder.layers.0.mlp.linear_fc1.bias", (1024,), True), + ("decoder.layers.0.mlp.experts.linear_fc1.weight", (64, 1024, 2048), False), + ("decoder.layers.0.mlp.experts.linear_fc1.bias", (64, 1024), False), + ("decoder.layers.0.mlp.experts.linear_fc2.weight", (64, 2048, 512), False), + ], +) +def test_glu_rechunk_only_applies_to_non_grouped_fc1_tensors(common, name, shape, expected): + assert common._requires_glu_rechunk(name, torch.empty(shape)) is expected + + +class _ParallelParam(torch.nn.Parameter): + def __new__(cls, data, partition_dim): + param = super().__new__(cls, data, requires_grad=False) + param.partition_dim = partition_dim + param.partition_stride = 1 + param.tensor_model_parallel = True + return param + + +class _CompletedWork: + def wait(self): + return None + + +def _patch_two_rank_all_gather(monkeypatch, common): + monkeypatch.setattr(common.mpu, "get_tensor_model_parallel_world_size", lambda: 2, raising=False) + monkeypatch.setattr(common.mpu, "get_tensor_model_parallel_group", lambda: None, raising=False) + monkeypatch.setattr(common.mpu, "get_expert_tensor_parallel_world_size", lambda: 2, raising=False) + monkeypatch.setattr(common.mpu, "get_expert_tensor_parallel_group", lambda: None, raising=False) + + def all_gather(outputs, _input, group, async_op=False): + del group + for rank, output in enumerate(outputs): + output.fill_(rank + 1) + return _CompletedWork() if async_op else None + + monkeypatch.setattr(common.dist, "all_gather", all_gather) + + +@pytest.mark.unit +def test_sync_all_gather_preserves_grouped_moe_expert_axis(monkeypatch, common): + _patch_two_rank_all_gather(monkeypatch, common) + param = _ParallelParam(torch.empty(2, 4, 3), partition_dim=2) + + gathered = common.all_gather_param("decoder.layers.0.mlp.experts.linear_fc1.weight", param) + + assert gathered.shape == (2, 4, 6) + assert torch.equal(gathered[..., :3], torch.ones(2, 4, 3)) + assert torch.equal(gathered[..., 3:], torch.full((2, 4, 3), 2.0)) + + +@pytest.mark.unit +def test_async_all_gather_preserves_grouped_moe_expert_axis(monkeypatch, common): + _patch_two_rank_all_gather(monkeypatch, common) + param = _ParallelParam(torch.empty(2, 4, 3), partition_dim=2) + info = SimpleNamespace(name="decoder.layers.0.mlp.experts.linear_fc1.weight") + + (gathered,) = common.all_gather_params_async([(info, param)]) + + assert gathered.shape == (2, 4, 6) + assert torch.equal(gathered[..., :3], torch.ones(2, 4, 3)) + assert torch.equal(gathered[..., 3:], torch.full((2, 4, 3), 2.0)) + + +@pytest.mark.unit +def test_sync_all_gather_keeps_2d_glu_rechunk_order(monkeypatch, common): + _patch_two_rank_all_gather(monkeypatch, common) + param = _ParallelParam(torch.empty(4, 3), partition_dim=0) + + gathered = common.all_gather_param("decoder.layers.0.mlp.linear_fc1.weight", param) + + expected = torch.cat( + [ + torch.ones(2, 3), + torch.full((2, 3), 2.0), + torch.ones(2, 3), + torch.full((2, 3), 2.0), + ] + ) + assert torch.equal(gathered, expected) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__]))