| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Test that the interaction between EPLB and FusedMoEFactory Layer is okay |
| 5 | |
| 6 | from dataclasses import dataclass |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | from vllm.config import VllmConfig, set_current_vllm_config |
| 12 | from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator |
| 13 | from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace |
| 14 | from vllm.distributed.parallel_state import ( |
| 15 | ensure_model_parallel_initialized, |
| 16 | get_eplb_group, |
| 17 | get_tp_group, |
| 18 | ) |
| 19 | from vllm.model_executor.layers.fused_moe import FusedMoEFactory, MoERunner |
| 20 | from vllm.platforms import current_platform |
| 21 | |
| 22 | from .eplb_utils import distributed_run, set_env_vars_and_device |
| 23 | |
| 24 | |
| 25 | @dataclass |
| 26 | class TestConfig: |
| 27 | num_layers: int |
| 28 | num_experts: int |
| 29 | num_local_experts: int |
| 30 | num_topk: int |
| 31 | hidden_size: int |
| 32 | intermediate_size: int |
| 33 | weight_dtype: torch.dtype |
| 34 | weight_scale_dtype: torch.dtype | None |
| 35 | column_major_scales: bool |
| 36 | |
| 37 | |
| 38 | def make_expert_weights( |
| 39 | layer_idx: int, |
| 40 | global_expert_idx: int, |
| 41 | global_num_experts: int, |
| 42 | tensor_shape: tuple[int, ...], |
| 43 | tensor_dtype: torch.dtype, |
| 44 | tensor_device: torch.device, |
| 45 | is_column_major: bool, |
| 46 | ) -> torch.Tensor: |
| 47 | assert len(tensor_shape) == 2 |
| 48 | |
| 49 | if is_column_major: |
| 50 | tensor_shape = (tensor_shape[1], tensor_shape[0]) |
| 51 | |
| 52 | x = torch.empty(tensor_shape, dtype=tensor_dtype, device=tensor_device) |
| 53 | value_offset = (layer_idx * global_num_experts + global_expert_idx) * x.numel() |
| 54 | x.view(-1).copy_( |
| 55 | torch.arange( |
| 56 | value_offset, |
| 57 | value_offset + x.numel(), |
| 58 | dtype=tensor_dtype, |
| 59 | device=tensor_device, |
| 60 | ) |
| 61 | ) |
| 62 | |
| 63 | if is_column_major: |
| 64 | x = torch.transpose(x, 1, 0) |
| 65 | assert not x.is_contiguous() |
| 66 | return x |
| 67 | |
| 68 | |
| 69 | def make_fused_moe_layer( |
| 70 | rank: int, |
| 71 | layer_idx: int, |
| 72 | test_config: TestConfig, |
| 73 | ) -> MoERunner: |
| 74 | fml = FusedMoEFactory( |
| 75 | num_experts=test_config.num_experts, |
| 76 | top_k=test_config.num_topk, |
| 77 | hidden_size=test_config.hidden_size, |
| 78 | intermediate_size=test_config.intermediate_size, |
| 79 | prefix=f"dummy_layer_{layer_idx}", |
| 80 | activation="silu", |
| 81 | params_dtype=test_config.weight_dtype, |
| 82 | ) |
| 83 | re = fml.routed_experts |
| 84 | |
| 85 | device = torch.device(f"{torch.accelerator.current_accelerator().type}:{rank}") |
| 86 | |
| 87 | from functools import partial |
| 88 | |
| 89 | _make_expert_weights = partial( |
| 90 | make_expert_weights, |
| 91 | layer_idx=layer_idx, |
| 92 | global_num_experts=test_config.num_experts, |
| 93 | tensor_device=device, |
| 94 | ) |
| 95 | |
| 96 | assert isinstance(re.w13_weight.data, torch.Tensor) |
| 97 | assert isinstance(re.w2_weight.data, torch.Tensor) |
| 98 | re.w13_weight.data = re.w13_weight.data.to(device=device) |
| 99 | re.w2_weight.data = re.w2_weight.data.to(device=device) |
| 100 | w13_weight = re.w13_weight.data |
| 101 | w2_weight = re.w2_weight.data |
| 102 | assert w13_weight.size(0) == test_config.num_local_experts |
| 103 | for i in range(test_config.num_local_experts): |
| 104 | g_i = rank * test_config.num_local_experts + i |
| 105 | w13_weight_e = w13_weight[i] |
| 106 | w2_weight_e = w2_weight[i] |
| 107 | w13_weight_e.copy_( |
| 108 | _make_expert_weights( |
| 109 | global_expert_idx=g_i, |
| 110 | tensor_shape=w13_weight_e.shape, |
| 111 | tensor_dtype=w13_weight_e.dtype, |
| 112 | is_column_major=False, |
| 113 | ) |
| 114 | ) |
| 115 | w2_weight_e.copy_( |
| 116 | _make_expert_weights( |
| 117 | global_expert_idx=g_i, |
| 118 | tensor_shape=w2_weight_e.shape, |
| 119 | tensor_dtype=w2_weight_e.dtype, |
| 120 | is_column_major=False, |
| 121 | ) |
| 122 | ) |
| 123 | |
| 124 | block_size = 16 |
| 125 | |
| 126 | def block_quant_scales_shape( |
| 127 | shape: tuple[int, ...], is_column_major: bool |
| 128 | ) -> tuple[int, ...]: |
| 129 | assert len(shape) == 3 |
| 130 | if not is_column_major: |
| 131 | return (shape[0], shape[1] // block_size, shape[2] // block_size) |
| 132 | else: |
| 133 | return (shape[0], shape[2] // block_size, shape[1] // block_size) |
| 134 | |
| 135 | is_column_major = test_config.column_major_scales |
| 136 | w13_weight_scale_inv = torch.empty( |
| 137 | block_quant_scales_shape(w13_weight.shape, is_column_major), |
| 138 | dtype=test_config.weight_dtype, |
| 139 | device=device, |
| 140 | ) |
| 141 | w2_weight_scale_inv = torch.empty( |
| 142 | block_quant_scales_shape(w2_weight.shape, is_column_major), |
| 143 | dtype=test_config.weight_dtype, |
| 144 | device=device, |
| 145 | ) |
| 146 | |
| 147 | for i in range(test_config.num_local_experts): |
| 148 | g_i = rank * test_config.num_local_experts + i |
| 149 | w13_s_e = w13_weight_scale_inv[i] |
| 150 | w2_s_e = w2_weight_scale_inv[i] |
| 151 | w13_s_e.copy_( |
| 152 | _make_expert_weights( |
| 153 | global_expert_idx=g_i, |
| 154 | tensor_shape=w13_s_e.shape, |
| 155 | tensor_dtype=w13_s_e.dtype, |
| 156 | # Fill data in row-major and then |
| 157 | # transpose if test_config requires col-major. |
| 158 | is_column_major=False, |
| 159 | ) |
| 160 | ) |
| 161 | w2_s_e.copy_( |
| 162 | _make_expert_weights( |
| 163 | global_expert_idx=g_i, |
| 164 | tensor_shape=w2_s_e.shape, |
| 165 | tensor_dtype=w2_s_e.dtype, |
| 166 | is_column_major=False, |
| 167 | ) |
| 168 | ) |
| 169 | if is_column_major: |
| 170 | w13_weight_scale_inv = torch.transpose(w13_weight_scale_inv, 1, 2) |
| 171 | w2_weight_scale_inv = torch.transpose(w2_weight_scale_inv, 1, 2) |
| 172 | assert not w13_weight_scale_inv.is_contiguous() |
| 173 | assert not w2_weight_scale_inv.is_contiguous() |
| 174 | |
| 175 | # Add scales to the parameter list |
| 176 | re.w13_weight_scale_inv = torch.nn.Parameter( |
| 177 | w13_weight_scale_inv, requires_grad=False |
| 178 | ) |
| 179 | re.w2_weight_scale_inv = torch.nn.Parameter( |
| 180 | w2_weight_scale_inv, requires_grad=False |
| 181 | ) |
| 182 | |
| 183 | return fml |
| 184 | |
| 185 | |
| 186 | def _test_eplb_fml(env, world_size: int, test_config: TestConfig): |
| 187 | # Initialize model parallel (using tensor parallel as an entrypoint |
| 188 | # to expert parallel) |
| 189 | set_env_vars_and_device(env) |
| 190 | |
| 191 | vllm_config = VllmConfig() |
| 192 | vllm_config.parallel_config.tensor_parallel_size = world_size |
| 193 | vllm_config.parallel_config.enable_expert_parallel = True |
| 194 | vllm_config.parallel_config.enable_eplb = True |
| 195 | |
| 196 | with set_current_vllm_config(vllm_config): |
| 197 | ensure_model_parallel_initialized( |
| 198 | tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1 |
| 199 | ) |
| 200 | |
| 201 | ep_group = get_tp_group().device_group |
| 202 | ep_rank = torch.distributed.get_rank() |
| 203 | |
| 204 | fml_layers = [ |
| 205 | make_fused_moe_layer(ep_rank, layer_idx, test_config) |
| 206 | for layer_idx in range(test_config.num_layers) |
| 207 | ] |
| 208 | rank_expert_weights = [fml.get_expert_weights() for fml in fml_layers] |
| 209 | |
| 210 | indices = torch.zeros( |
| 211 | test_config.num_layers, test_config.num_experts, dtype=torch.long |
| 212 | ) |
| 213 | for lidx in range(test_config.num_layers): |
| 214 | indices[lidx] = torch.Tensor(range(test_config.num_experts)) |
| 215 | |
| 216 | shuffled_indices = torch.zeros_like(indices) |
| 217 | for lidx in range(test_config.num_layers): |
| 218 | shuffled_indices[lidx] = torch.randperm(test_config.num_experts) |
| 219 | |
| 220 | expert_buffer = [torch.empty_like(w) for w in rank_expert_weights[0]] |
| 221 | default_backend = "torch_xccl" if current_platform.is_xpu() else "torch_nccl" |
| 222 | communicator = create_eplb_communicator( |
| 223 | group_coordinator=get_eplb_group(), |
| 224 | backend=default_backend, |
| 225 | expert_weights=rank_expert_weights, |
| 226 | expert_buffer=expert_buffer, |
| 227 | ) |
| 228 | rearrange_expert_weights_inplace( |
| 229 | indices, |
| 230 | shuffled_indices, |
| 231 | rank_expert_weights, |
| 232 | expert_buffer, |
| 233 | ep_group, |
| 234 | communicator, |
| 235 | ) |
| 236 | |
| 237 | num_local_experts = test_config.num_local_experts |
| 238 | num_global_experts = test_config.num_experts |
| 239 | for lidx, fml in enumerate(fml_layers): |
| 240 | for name, w in fml.named_parameters(): |
| 241 | for e in range(num_local_experts): |
| 242 | g_e = shuffled_indices[lidx][ep_rank * num_local_experts + e] |
| 243 | ref = make_expert_weights( |
| 244 | layer_idx=lidx, |
| 245 | global_expert_idx=int(g_e.item()), |
| 246 | global_num_experts=num_global_experts, |
| 247 | tensor_shape=w[e].shape, |
| 248 | tensor_dtype=w[e].dtype, |
| 249 | tensor_device=w[e].device, |
| 250 | is_column_major=not w[e].is_contiguous(), |
| 251 | ) |
| 252 | assert w[e].shape == ref.shape and w[e].stride() == ref.stride(), ( |
| 253 | f"w[{e}] {w[e].size()} {w[e].stride()} vs " |
| 254 | f"ref {ref.size()} {ref.stride()}" |
| 255 | ) |
| 256 | torch.testing.assert_close(w[e], ref) |
| 257 | |
| 258 | |
| 259 | @pytest.mark.parametrize("world_size", [2]) |
| 260 | @pytest.mark.parametrize("num_layers", [4]) |
| 261 | @pytest.mark.parametrize("num_experts", [16]) |
| 262 | @pytest.mark.parametrize("hidden_size", [256]) |
| 263 | @pytest.mark.parametrize("intermediate_size", [256]) |
| 264 | @pytest.mark.parametrize("column_major_scales", [True, False]) |
| 265 | def test_eplb_fml( |
| 266 | world_size: int, |
| 267 | num_layers: int, |
| 268 | num_experts: int, |
| 269 | hidden_size: int, |
| 270 | intermediate_size: int, |
| 271 | column_major_scales: bool, |
| 272 | ): |
| 273 | if torch.accelerator.device_count() < world_size: |
| 274 | pytest.skip(f"Need at least {world_size} GPUs to run the test") |
| 275 | |
| 276 | num_local_experts = num_experts // world_size |
| 277 | num_topk = 4 |
| 278 | # The dtypes are fine as we are essentially just checking data-copies |
| 279 | weight_dtype = torch.bfloat16 |
| 280 | weight_scale_dtype = torch.bfloat16 |
| 281 | |
| 282 | test_config = TestConfig( |
| 283 | num_layers=num_layers, |
| 284 | num_experts=num_experts, |
| 285 | num_local_experts=num_local_experts, |
| 286 | num_topk=num_topk, |
| 287 | hidden_size=hidden_size, |
| 288 | intermediate_size=intermediate_size, |
| 289 | weight_dtype=weight_dtype, |
| 290 | weight_scale_dtype=weight_scale_dtype, |
| 291 | column_major_scales=column_major_scales, |
| 292 | ) |
| 293 | |
| 294 | distributed_run( |
| 295 | _test_eplb_fml, |
| 296 | world_size, |
| 297 | test_config, |
| 298 | ) |