deepspeedai--deepspeed/tests/unit/v1/moe/test_autoep_autotp_runtime.py

# Copyright (c) DeepSpeed Team.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team
"""Runtime wiring checks for AutoEP + AutoTP folding."""

import torch
import torch.nn as nn

import deepspeed
import deepspeed.comm as dist
from deepspeed.module_inject.auto_ep_config import AutoEPConfig, MoELayerSpec
from deepspeed.module_inject.auto_ep_folding import FoldingGroupHandles, build_folding_spec, local_folding_ranks
from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer
from deepspeed.module_inject.auto_tp import AutoTP
from deepspeed.utils import groups
from unit.common import DistributedTest
from unit.v1.moe.autoep_test_utils import (
    MockMoEOnlyTransformer,
    engine_input_dtype,
    make_autoep_config,
    run_cpu_gloo_test,
    seed_everything,
    skip_unless_h100_tests_enabled,
)


def _make_spec(**overrides):
    defaults = dict(
        moe_module_name="model.layers.0.mlp",
        model_family="mixtral",
        router_name="gate",
        experts_name="experts",
        expert_storage="fused_3d",
        expert_w1_name="gate_up_proj",
        expert_w2_name="down_proj",
        expert_w3_name=None,
        num_experts=4,
        top_k=2,
        hidden_size=8,
        ffn_hidden_size=16,
        score_func="softmax",
        score_apply="post",
        route_norm=True,
        gate_bias=False,
        return_router_logits=False,
        router_logits_capture_target="none",
        router_logits_capture_index=None,
        router_logits_capture_layer_name=None,
        has_shared_experts=False,
        shared_experts_name="",
        shared_experts_gate_name="",
    )
    defaults.update(overrides)
    return MoELayerSpec(**defaults)


class TinySourceMoE(nn.Module):

    def __init__(self):
        super().__init__()
        self.gate = nn.Linear(8, 4, bias=False)
        self.experts = nn.Module()
        self.experts.gate_up_proj = nn.Parameter(torch.randn(4, 32, 8))
        self.experts.down_proj = nn.Parameter(torch.randn(4, 8, 16))


def test_folded_layer_binds_explicit_group_handles(monkeypatch):
    layer = AutoEPMoELayer(_make_spec(), TinySourceMoE(), ep_size=2, ep_rank=0, config=AutoEPConfig(enabled=True))
    spec = build_folding_spec(world_size=4, pp_size=1, tp_size=2, ep_size=2, etp_size=1)
    local = local_folding_ranks(0, spec)
    handles = FoldingGroupHandles(
        spec=spec,
        tp_group=object(),
        dense_dp_group=object(),
        ep_group=object(),
        edp_group=object(),
        ep_group_name="ep_size_2",
        tp_ranks=local["tp"],
        dense_dp_ranks=local["dense_dp"],
        ep_ranks=local["ep"],
        edp_ranks=local["edp"],
    )
    monkeypatch.setattr("deepspeed.module_inject.auto_ep_layer.dist.get_rank", lambda group=None: 0)

    layer.set_deepspeed_parallelism(folding_group_handles=handles)

    assert layer.folding_group_handles is handles
    assert layer.tp_group is handles.tp_group
    assert layer.ep_group is handles.ep_group
    assert layer.ep_group_name == "ep_size_2"


def test_autotp_reaches_autoep_shared_experts(monkeypatch):

    class AutoEPLike(nn.Module):

        def __init__(self):
            super().__init__()
            self._is_autoep_layer = True
            self.shared_experts = nn.Linear(8, 8, bias=False)
            self.shared_experts_gate = nn.Linear(8, 8, bias=False)

    model = nn.Module()
    model.moe = AutoEPLike()
    autotp = AutoTP.__new__(AutoTP)
    calls = []
    monkeypatch.setattr(autotp, "_replace_autoep_shared_experts", lambda child, name: calls.append((child, name)))

    AutoTP._replace_module(autotp, model)

    assert calls == [(model.moe, "moe")]


def _folded_config(zero_stage=0, *, ep_size=2, mixed_precision=True):
    config = make_autoep_config(zero_stage=zero_stage, ep_size=ep_size, mixed_precision=mixed_precision)
    if not mixed_precision:
        config["optimizer"]["params"]["torch_adam"] = True
    config["tensor_parallel"] = {
        "autotp_size": 2,
        "partition_config": {
            "use_default_specs": False,
            "layer_specs": [{
                "patterns": [".*\\.weight$"],
                "partition_type": "skip",
            }],
        },
    }
    return config


def _tp_consistent_input(engine, *, seed=1234):
    torch.manual_seed(seed)
    x = torch.randn(1, 4, 64, device=engine.device, dtype=engine_input_dtype(engine))
    dist.broadcast(x, groups.get_tensor_model_parallel_src_rank(), group=groups.get_tensor_model_parallel_group())
    return x


def _initialize_folded_engine(*, zero_stage=0, ep_size=2, mixed_precision=True):
    seed_everything(1234)
    return deepspeed.initialize(model=MockMoEOnlyTransformer(),
                                config=_folded_config(zero_stage=zero_stage,
                                                      ep_size=ep_size,
                                                      mixed_precision=mixed_precision))


def _assert_nonzero_named_grad(engine, *name_fragments):
    grad_total = 0.0
    matched = False
    for name, param in engine.module.named_parameters():
        if not any(fragment in name for fragment in name_fragments):
            continue
        if param.grad is None:
            continue
        matched = True
        grad_total += param.grad.detach().float().abs().sum().item()
    assert matched, f"no gradients found for parameters matching {name_fragments}"
    assert grad_total > 0.0


def _cpu_folded_runtime_worker(_rank, _world_size, _shared_tmpdir):
    engine, _, _, _ = _initialize_folded_engine(zero_stage=0, mixed_precision=False)
    assert engine.autotp_size() == 2
    assert groups.get_tensor_model_parallel_world_size() == 2
    folded_layers = [module for module in engine.module.modules() if isinstance(module, AutoEPMoELayer)]
    assert folded_layers
    assert all(layer.folding_group_handles is not None for layer in folded_layers)

    x = _tp_consistent_input(engine)
    loss = engine(x).float().mean()
    engine.backward(loss)
    _assert_nonzero_named_grad(engine, "experts.")
    _assert_nonzero_named_grad(engine, "router", "gate")
    engine.step()
    assert torch.isfinite(loss.detach()).item()


def test_cpu_gloo_folded_runtime_smoke(tmpdir):
    run_cpu_gloo_test(_cpu_folded_runtime_worker, tmpdir, world_size=4)


class TestH100FoldedRuntime(DistributedTest):
    world_size = 4
    reuse_dist_env = False

    def test_h100_folded_tp2_ep2_runtime(self):
        skip_unless_h100_tests_enabled("H100 runtime node")

        engine, _, _, _ = _initialize_folded_engine(zero_stage=0)
        assert engine.autotp_size() == 2
        assert groups.get_tensor_model_parallel_world_size() == 2
        folded_layers = [module for module in engine.module.modules() if isinstance(module, AutoEPMoELayer)]
        assert folded_layers
        assert all(layer.folding_group_handles is not None for layer in folded_layers)

        x = _tp_consistent_input(engine)
        loss = engine(x).float().mean()
        engine.backward(loss)
        _assert_nonzero_named_grad(engine, "experts.")
        _assert_nonzero_named_grad(engine, "router", "gate")
        engine.step()
        assert torch.isfinite(loss.detach()).item()


class TestH100FoldedRuntimeReference(DistributedTest):
    world_size = 4
    reuse_dist_env = False

    def test_h100_folded_tp2_ep2_finite_loss_smoke(self):
        skip_unless_h100_tests_enabled("H100 benchmark node")

        engine, _, _, _ = _initialize_folded_engine(zero_stage=0)
        x = _tp_consistent_input(engine)
        losses = []
        for _ in range(2):
            loss = engine(x).float().mean()
            engine.backward(loss)
            engine.step()
            losses.append(float(loss.detach().cpu()))
        assert all(torch.isfinite(torch.tensor(value)) for value in losses)


class TestH100FoldedRuntimeTP2EP4(DistributedTest):
    world_size = 8
    reuse_dist_env = False

    def test_h100_folded_tp2_ep4_runtime(self):
        skip_unless_h100_tests_enabled("H100 TP2-EP4 runtime node")

        engine, _, _, _ = _initialize_folded_engine(zero_stage=0, ep_size=4)
        assert engine.autotp_size() == 2
        assert groups.get_tensor_model_parallel_world_size() == 2
        folded_layers = [module for module in engine.module.modules() if isinstance(module, AutoEPMoELayer)]
        assert folded_layers
        assert all(layer.folding_group_handles is not None for layer in folded_layers)

        x = _tp_consistent_input(engine)
        loss = engine(x).float().mean()
        engine.backward(loss)
        _assert_nonzero_named_grad(engine, "experts.")
        _assert_nonzero_named_grad(engine, "router", "gate")
        engine.step()
        assert torch.isfinite(loss.detach()).item()