# Copyright (c) Microsoft Corporation. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team import torch import deepspeed import pytest import gc import random from unit.common import DistributedTest from unit.simple_model import SimplePRMoEModel, SimpleMoEModel, sequence_dataloader import deepspeed.comm as dist import deepspeed.moe.sharded_moe as sharded_moe from deepspeed import get_accelerator from deepspeed.moe.layer import MoE from deepspeed.moe.sharded_moe import top1gating, top2gating, topkgating from deepspeed.moe.utils import split_params_into_different_moe_groups_for_optimizer, is_moe_param from deepspeed.utils.torch import required_torch_version @pytest.mark.parametrize("zero_stage", [0, 1, 2]) class TestSimpleMoE(DistributedTest): world_size = 2 def test(self, zero_stage): if not required_torch_version(min_version=1.8): pytest.skip("DeepSpeed MoE tests need torch 1.8 or higher to run correctly") config_dict = { "train_micro_batch_size_per_gpu": 1, "steps_per_print": 1, "optimizer": { "type": "Adam", "params": { "lr": 0.00015 } }, "fp16": { "enabled": True }, "zero_optimization": { "stage": zero_stage } } # should automatically create moe param groups in deepspeed backend hidden_dim = 16 model = SimpleMoEModel(hidden_dim=hidden_dim, ep_size=1) model, optimizer, _, _ = deepspeed.initialize(config=config_dict, model=model) data_loader = sequence_dataloader(model=model, total_samples=50, hidden_dim=hidden_dim, device=model.device, dtype=torch.float16) for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) model.backward(loss) model.step() @pytest.mark.parametrize("ep_size", [2, 4]) @pytest.mark.parametrize("zero_stage", [0, 1, 2]) @pytest.mark.parametrize("use_residual", [True, False]) class TestMoE(DistributedTest): world_size = 4 def test(self, ep_size, zero_stage, use_residual): if not required_torch_version(min_version=1.8): pytest.skip("DeepSpeed MoE tests need torch 1.8 or higher to run correctly") config_dict = { "train_micro_batch_size_per_gpu": 1, "steps_per_print": 1, "fp16": { "enabled": True }, "zero_optimization": { "stage": zero_stage } } hidden_dim = 16 # E+D -- ep_size = 2 # E only -- ep_size = 4 model = SimpleMoEModel(hidden_dim, ep_size=ep_size, use_residual=use_residual) param_group = {'params': [p for p in model.parameters()], 'name': 'random-unique-name'} params = split_params_into_different_moe_groups_for_optimizer(param_group) optimizer = torch.optim.AdamW(params=params) model, optimizer, _, _ = deepspeed.initialize(config=config_dict, model=model, optimizer=optimizer, dist_init_required=False) #dist_init_required=False -- parameterize to True/False? data_loader = sequence_dataloader(model=model, total_samples=50, hidden_dim=hidden_dim, device=model.device, dtype=torch.float16) def strict_average_tensor(tensor, communication_data_type: torch.dtype): process_group = optimizer.dp_process_group curr_size = 0 pg_offsets = [] ipg_bucket = optimizer.ipg_buckets[communication_data_type] for i, param_idx, param_id in ipg_bucket.params: param = optimizer.bit16_groups[i][param_idx] process_group = optimizer.dp_process_group if ipg_bucket.has_moe_params: process_group = optimizer.expert_dp_process_group[param.group_name] if is_moe_param( param) else optimizer.dp_process_group partition_ids = optimizer.param_to_partition_ids[i][param_id] # Get all partition ids + their offsets partition_offsets = [] for partition_id in partition_ids: offset = optimizer.grad_start_offset[i][partition_id][param_id] partition_offsets.append(offset) partition_offsets.sort() # Calculate rank and offsets for grad slices for idx, offset in enumerate(partition_offsets): # Calculate numel for grad slice depending on partition location if idx == len(partition_offsets) - 1: # Last partition_id uses its own offset numel = param.numel() - offset else: # Set numel to next partition's offset numel = partition_offsets[idx + 1] - offset pg_offsets.append((curr_size, process_group)) curr_size += numel def strict_narrow(dim, start, length): lo, hi = 0, len(pg_offsets) - 1 while lo < hi: mi = lo + (hi - lo) // 2 if pg_offsets[mi][0] >= start: hi = mi else: lo = mi + 1 curr_slice, reduce_process_group = lo, pg_offsets[lo][1] while curr_slice < len(pg_offsets) and start + length > pg_offsets[curr_slice][0]: assert reduce_process_group == pg_offsets[curr_slice][ 1], "reduce process_group does not match the parameter's process_group" curr_slice += 1 return orig_narrow(dim, start, length) # real call orig_narrow, tensor.narrow = tensor.narrow, strict_narrow type(optimizer).average_tensor(optimizer, tensor, communication_data_type) # real call tensor.narrow = orig_narrow if "average_tensor" in dir(optimizer): optimizer.average_tensor = strict_average_tensor for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) model.backward(loss) model.step() gc.collect() # Must do this or we get a memory leak in this test @pytest.mark.parametrize("ep_size, use_residual", [(2, True), (2, False)]) class TestPRMoE(DistributedTest): world_size = 4 def test(self, ep_size, use_residual): if not required_torch_version(min_version=1.8): pytest.skip("DeepSpeed MoE tests need torch 1.8 or higher to run correctly") config_dict = {"train_batch_size": 8, "steps_per_print": 1, "fp16": {"enabled": True}} hidden_dim = 16 # E+D -- ep_size = 2 # E only -- ep_size = 4 model = SimplePRMoEModel(hidden_dim, ep_size=ep_size, use_residual=use_residual) optimizer = torch.optim.AdamW(params=model.parameters()) model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, optimizer=optimizer, dist_init_required=False) data_loader = sequence_dataloader(model=model, total_samples=50, hidden_dim=hidden_dim, device=model.device, dtype=torch.float16) for n, batch in enumerate(data_loader): loss = model(batch[0], batch[1]) model.backward(loss) model.step() class TestTopk(DistributedTest): world_size = 2 def test(self): device = get_accelerator().current_device_name() if dist.get_rank() == 0: logits = torch.rand(2, 2, device=device) elif dist.get_rank() == 1: logits = torch.rand(10, 2, device=device) output = top1gating(logits=logits, capacity_factor=1, min_capacity=0, used_token=None, noisy_gate_policy=None, drop_tokens=False, use_rts=True, use_tutel=False) class TestMoESingleton(DistributedTest): world_size = 2 @pytest.mark.parametrize("ep_size, expected_calls", [(1, 0), (2, 2)], ids=["single", "multi"]) def test_all_to_all(self, monkeypatch, ep_size, expected_calls): if not required_torch_version(min_version=1.8): pytest.skip("DeepSpeed MoE tests need torch 1.8 or higher to run correctly") config_dict = {"train_micro_batch_size_per_gpu": 1, "steps_per_print": 1} hidden_dim = 8 expert = torch.nn.Sequential(torch.nn.Linear(hidden_dim, hidden_dim), torch.nn.Linear(hidden_dim, hidden_dim)) model = MoE(hidden_size=hidden_dim, expert=expert, num_experts=2, ep_size=ep_size, k=1, min_capacity=0) optimizer = torch.optim.AdamW(params=model.parameters()) model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, optimizer=optimizer, dist_init_required=False) all_to_all_calls = [] def counted_all_to_all(group, input): all_to_all_calls.append((group, input.shape)) return input monkeypatch.setattr(sharded_moe._AllToAll, "apply", counted_all_to_all) x = torch.randn(1, 4, hidden_dim, device=model.device, requires_grad=True) output, l_aux, _ = model(x) assert len(all_to_all_calls) == expected_calls loss = output.float().sum() + l_aux.float() model.backward(loss) assert len(all_to_all_calls) == expected_calls assert x.grad is not None assert any(param.grad is not None for param in model.module.parameters()) @pytest.mark.parametrize("gate_fn, capacity_args", [(top1gating, (1, 0)), (top2gating, (1, 0)), (topkgating, (3, 1, 0))], ids=["top1", "top2", "topk"]) @pytest.mark.parametrize("ep_world_size, expected_calls", [(1, 0), (2, 1)], ids=["single", "multi"]) def test_capacity(self, monkeypatch, gate_fn, capacity_args, ep_world_size, expected_calls): if not required_torch_version(min_version=1.8): pytest.skip("DeepSpeed MoE tests need torch 1.8 or higher to run correctly") ep_group = None if ep_world_size == 1: for rank in range(dist.get_world_size()): group = dist.new_group([rank]) if rank == dist.get_rank(): ep_group = group else: ep_group = dist.new_group(list(range(dist.get_world_size()))) all_reduce_calls = [] original_all_reduce = sharded_moe.dist.all_reduce def counted_all_reduce(tensor, op=dist.ReduceOp.SUM, group=None): all_reduce_calls.append((tensor, op, group)) return original_all_reduce(tensor, op=op, group=group) monkeypatch.setattr(sharded_moe.dist, "all_reduce", counted_all_reduce) device = get_accelerator().current_device_name() logits = torch.randn(8, 4, device=device) gate_fn(logits, *capacity_args, drop_tokens=False, ep_group=ep_group) assert len(all_reduce_calls) == expected_calls if all_reduce_calls: _, op, group = all_reduce_calls[0] assert op == dist.ReduceOp.MAX assert group is ep_group def test_no_ep_group(self, monkeypatch): if not required_torch_version(min_version=1.8): pytest.skip("DeepSpeed MoE tests need torch 1.8 or higher to run correctly") def fail_collective(*args, **kwargs): raise AssertionError("ep_group=None should not enter expert-parallel collective code") monkeypatch.setattr(sharded_moe.dist, "get_world_size", fail_collective) monkeypatch.setattr(sharded_moe.dist, "all_reduce", fail_collective) device = get_accelerator().current_device_name() logits = torch.randn(8, 4, device=device) top2gating(logits, 1, 0, drop_tokens=False, ep_group=None, top2_2nd_expert_sampling=False) class TestTopkGate(DistributedTest): def test(self): def check_equal(logits, cap, sparse_truth, res): m, n = logits.shape dispatch_mask_truth = torch.zeros(m, n, cap) i, j, k = sparse_truth.t() dispatch_mask_truth[i, j, k] = 1 assert (torch.equal(dispatch_mask_truth, res)) #s=4 e=4 topk=2 cap=2(s*topk/e) logits = torch.tensor([[0.11, 0.2, 0.1, 0.3], [0.3, 0.4, 0.11, 0.1], [0.11, 0.1, 0.6, 0.5], [0.1, 0.11, 0.7, 0.8]]) logits *= dist.get_rank() + 1 probs_dispatch_res = topkgating(logits, 2, 1, min_capacity=1, drop_policy='probs')[2] probs_sec_sparse = torch.tensor([[0, 1, 0], [1, 0, 0], [1, 1, 1], [2, 2, 0], [2, 3, 0], [3, 2, 1], [3, 3, 1]]) check_equal(logits, 2, probs_sec_sparse, probs_dispatch_res) position_sec_sparse = torch.tensor([[0, 1, 0], [0, 3, 0], [1, 0, 0], [1, 1, 1], [2, 2, 0], [2, 3, 1], [3, 2, 1]]) position_dispatch_res = topkgating(logits, 2, 1, min_capacity=1, drop_policy='position')[2] check_equal(logits, 2, position_sec_sparse, position_dispatch_res) #s=4 e=6 topk=3 cap=2(s*topk/e) logits2 = torch.tensor([[0.5858, 0.4801, 0.6269, 0.5397, 0.9722, 0.7034], [0.5445, 0.6332, 0.4519, 0.6308, 0.0519, 0.6450], [0.4874, 0.8110, 0.7467, 0.8474, 0.0277, 0.3068], [0.8570, 0.6714, 0.5310, 0.3274, 0.4836, 0.9892]]) logits2 *= dist.get_rank() + 1 #top3 full mask #prob_mask #postion_mask #0 0 1 0 1 1 #0 0 1 0 1 0 #0 0 1 0 1 1 #0 1 0 1 0 1 #0 1 0 1 0 1 #0 1 0 1 0 1 #0 1 1 1 0 0 #0 1 1 1 0 0 #0 1 1 1 0 0 #1 1 0 0 0 1 #1 0 0 0 0 1 #1 0 0 0 0 0 probs_dispatch_res = topkgating(logits2, 3, 1, min_capacity=1, drop_policy='probs')[2] probs_sec_sparse = torch.tensor([[0, 2, 0], [0, 4, 0], [1, 1, 0], [1, 3, 0], [1, 5, 0], [2, 1, 1], [2, 2, 1], [2, 3, 1], [3, 0, 0], [3, 5, 1]]) check_equal(logits2, 2, probs_sec_sparse, probs_dispatch_res) position_sec_sparse = torch.tensor([[0, 2, 0], [0, 4, 0], [0, 5, 0], [1, 1, 0], [1, 3, 0], [1, 5, 1], [2, 1, 1], [2, 2, 1], [2, 3, 1], [3, 0, 0]]) position_dispatch_res = topkgating(logits2, 3, 1, min_capacity=1, drop_policy='position')[2] check_equal(logits2, 2, position_sec_sparse, position_dispatch_res) #s=4 e=4 topk=2 drop_tokens=False logits3 = torch.tensor([[0.95, 0.85, 0.90, 0.80], [0.70, 0.65, 0.75, 0.60], [0.50, 0.55, 0.45, 0.40], [0.35, 0.30, 0.25, 0.20]]) logits3 *= dist.get_rank() + 1 dispatch_res = topkgating(logits3, 2, 1, min_capacity=1, drop_tokens=False)[2] sec_sparse = torch.tensor([[0, 0, 0], [0, 2, 0], [1, 0, 1], [1, 2, 1], [2, 0, 2], [2, 1, 0], [3, 0, 3], [3, 1, 1]]) check_equal(logits3, 4, sec_sparse, dispatch_res) class TestExpertWeightGradWithZero(DistributedTest): world_size = 2 @pytest.mark.parametrize("zero_stage", [0, 1, 2]) def test(self, zero_stage): if not required_torch_version(min_version=1.8): pytest.skip("DeepSpeed MoE tests need torch 1.8 or higher to run correctly") def seed_everything(seed=11): random.seed(seed) torch.manual_seed(seed) get_accelerator().manual_seed(seed) get_accelerator().manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False def get_state_dict_ep2(state_dict): """ convert state_dict from EP=1 to EP=2 """ rank = int(deepspeed.comm.get_rank()) ep_state_dict = dict() dst_sub_key = "deepspeed_moe.experts.deepspeed_experts.0" src_sub_key = f"deepspeed_moe.experts.deepspeed_experts.{rank}" for moe_layer in ["moe_1", "moe_2"]: for mlp_in_moe in [0, 1]: dst_key = f"{moe_layer}.{dst_sub_key}.{mlp_in_moe}" src_key = f"{moe_layer}.{src_sub_key}.{mlp_in_moe}" ep_state_dict[f"{dst_key}.weight"] = state_dict[f"{src_key}.weight"].detach().clone() ep_state_dict[f"{dst_key}.bias"] = state_dict[f"{src_key}.bias"].detach().clone() for key in state_dict.keys(): if "deepspeed_moe.experts.deepspeed_experts" not in key: ep_state_dict[key] = state_dict[key].detach().clone() return ep_state_dict def get_models(hidden_dim): model_ep1 = SimpleMoEModel(hidden_dim=hidden_dim, num_experts=2, ep_size=1, use_rts=False) model_ep2 = SimpleMoEModel(hidden_dim=hidden_dim, num_experts=2, ep_size=2, use_rts=False) state_dict_ep1 = model_ep1.state_dict() state_dict_ep2 = get_state_dict_ep2(state_dict_ep1) model_ep2.load_state_dict(state_dict_ep2) model_ep1, _, _, _ = deepspeed.initialize(config=config_dict, model=model_ep1) model_ep2, _, _, _ = deepspeed.initialize(config=config_dict, model=model_ep2) return model_ep1, model_ep2 def extract_expert_grad(model, expert_id): def _get_weight_bias(experts): return ([deepspeed.utils.safe_get_full_grad(expert[0].weight) for expert in experts][expert_id].detach().clone(), [deepspeed.utils.safe_get_full_grad(expert[0].bias) for expert in experts][expert_id].detach().clone(), [deepspeed.utils.safe_get_full_grad(expert[1].weight) for expert in experts][expert_id].detach().clone(), [deepspeed.utils.safe_get_full_grad(expert[1].bias) for expert in experts][expert_id].detach().clone()) return (*_get_weight_bias(model.moe_1.deepspeed_moe.experts.deepspeed_experts), *_get_weight_bias(model.moe_2.deepspeed_moe.experts.deepspeed_experts)) seed_everything() config_dict = { "train_micro_batch_size_per_gpu": 1, "steps_per_print": 1, "optimizer": { "type": "Adam", "params": { "lr": 0.1, } }, "zero_optimization": { "stage": zero_stage } } hidden_dim = 4 total_samples = 2 rank = deepspeed.comm.get_rank() model_ep1, model_ep2 = get_models(hidden_dim) data_loader = sequence_dataloader(model=model_ep1, total_samples=total_samples, hidden_dim=hidden_dim, device=model_ep1.device, dtype=torch.float32) expert_weight_grad_ep1 = [] expert_weight_grad_ep2 = [] for batch in data_loader: loss_ep1 = model_ep1(batch[0], batch[1]) loss_ep2 = model_ep2(batch[0], batch[1]) model_ep1.backward(loss_ep1) model_ep2.backward(loss_ep2) expert_weight_grad_ep1.extend(extract_expert_grad(model_ep1, rank)) expert_weight_grad_ep2.extend(extract_expert_grad(model_ep2, 0)) model_ep1.step() model_ep2.step() assert len(expert_weight_grad_ep1) == len(expert_weight_grad_ep2) for grad_from_ep1, grad_from_ep2 in zip(expert_weight_grad_ep1, expert_weight_grad_ep2): assert torch.allclose(grad_from_ep1, grad_from_ep2, atol=0, rtol=1e-4)