fix(update_weight): preserve grouped MoE expert axis during GLU rechunk

.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"

.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},

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.

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__]))