| 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 MoERunner Layer is okay for DP w/ NVFP4 |
| 5 | |
| 6 | from dataclasses import dataclass |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | from tests.kernels.moe.utils import make_test_quant_config |
| 12 | from vllm.config import VllmConfig, set_current_vllm_config |
| 13 | from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator |
| 14 | from vllm.distributed.eplb.eplb_state import EplbLayerState |
| 15 | from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace |
| 16 | from vllm.distributed.parallel_state import ( |
| 17 | ensure_model_parallel_initialized, |
| 18 | get_dp_group, |
| 19 | get_eplb_group, |
| 20 | ) |
| 21 | from vllm.forward_context import set_forward_context |
| 22 | from vllm.model_executor.layers.fused_moe.layer import FusedMoEFactory, MoERunner |
| 23 | from vllm.model_executor.layers.quantization.modelopt import ( |
| 24 | ModelOptNvFp4Config, |
| 25 | ModelOptNvFp4FusedMoE, |
| 26 | ) |
| 27 | from vllm.platforms import current_platform |
| 28 | |
| 29 | from .eplb_utils import distributed_run, set_env_vars_and_device |
| 30 | |
| 31 | |
| 32 | @dataclass |
| 33 | class TestConfig: |
| 34 | num_layers: int |
| 35 | num_experts: int |
| 36 | num_local_experts: int |
| 37 | num_topk: int |
| 38 | hidden_size: int |
| 39 | intermediate_size: int |
| 40 | num_tokens: int |
| 41 | moe_backend: str |
| 42 | |
| 43 | |
| 44 | def make_fused_moe_layer( |
| 45 | rank: int, |
| 46 | layer_idx: int, |
| 47 | test_config: TestConfig, |
| 48 | ) -> MoERunner: |
| 49 | quant_config = None |
| 50 | |
| 51 | device = torch.device(f"cuda:{rank}") |
| 52 | |
| 53 | quant_config = ModelOptNvFp4Config( |
| 54 | is_checkpoint_nvfp4_serialized=True, |
| 55 | kv_cache_quant_algo=None, |
| 56 | exclude_modules=[], |
| 57 | ) |
| 58 | |
| 59 | fml = FusedMoEFactory( |
| 60 | num_experts=test_config.num_experts, |
| 61 | top_k=test_config.num_topk, |
| 62 | hidden_size=test_config.hidden_size, |
| 63 | intermediate_size=test_config.intermediate_size, |
| 64 | prefix=f"dummy_layer_{layer_idx}", |
| 65 | activation="silu", |
| 66 | params_dtype=torch.bfloat16, |
| 67 | quant_config=quant_config, |
| 68 | ) |
| 69 | |
| 70 | nvfp4_fused_moe = ModelOptNvFp4FusedMoE(quant_config, fml) |
| 71 | nvfp4_fused_moe.create_weights( |
| 72 | fml, |
| 73 | test_config.num_local_experts, |
| 74 | test_config.hidden_size, |
| 75 | test_config.intermediate_size, |
| 76 | params_dtype=torch.uint8, |
| 77 | global_num_experts=test_config.num_experts, |
| 78 | ) |
| 79 | |
| 80 | fml = fml.to(device) |
| 81 | re = fml.routed_experts |
| 82 | w1_q, w2_q, quant_config = make_test_quant_config( |
| 83 | test_config.num_local_experts, |
| 84 | test_config.intermediate_size, |
| 85 | test_config.hidden_size, |
| 86 | in_dtype=torch.bfloat16, |
| 87 | quant_dtype="nvfp4", |
| 88 | block_shape=None, |
| 89 | per_act_token_quant=False, |
| 90 | ) |
| 91 | |
| 92 | re.w13_weight.data = w1_q |
| 93 | re.w2_weight.data = w2_q |
| 94 | |
| 95 | re.w2_input_scale.data = torch.randn_like(re.w2_input_scale.data) / 5 |
| 96 | re.w13_input_scale.data = torch.randn_like(re.w13_input_scale.data) / 5 |
| 97 | re.w2_weight_scale_2.data = torch.randn_like(re.w2_weight_scale_2.data) / 5 |
| 98 | re.w13_weight_scale_2.data = torch.randn_like(re.w13_weight_scale_2.data) / 5 |
| 99 | re.w2_weight_scale.data = ( |
| 100 | torch.randn(re.w2_weight_scale.data.shape, device=device) / 5 |
| 101 | ).to(re.w2_weight_scale.data.dtype) |
| 102 | re.w13_weight_scale.data = ( |
| 103 | torch.randn(re.w13_weight_scale.data.shape, device=device) / 5 |
| 104 | ).to(re.w13_weight_scale.data.dtype) |
| 105 | |
| 106 | nvfp4_fused_moe.process_weights_after_loading(re) |
| 107 | |
| 108 | fml.maybe_init_modular_kernel() |
| 109 | |
| 110 | return fml |
| 111 | |
| 112 | |
| 113 | def _test_eplb_fml(env, world_size: int, test_config: TestConfig): |
| 114 | set_env_vars_and_device(env) |
| 115 | |
| 116 | vllm_config = VllmConfig() |
| 117 | vllm_config.parallel_config.data_parallel_size = world_size |
| 118 | vllm_config.parallel_config.enable_expert_parallel = True |
| 119 | vllm_config.kernel_config.moe_backend = test_config.moe_backend |
| 120 | |
| 121 | with set_current_vllm_config(vllm_config): |
| 122 | ensure_model_parallel_initialized( |
| 123 | tensor_model_parallel_size=1, pipeline_model_parallel_size=1 |
| 124 | ) |
| 125 | |
| 126 | ep_group = get_dp_group().cpu_group |
| 127 | ep_rank = torch.distributed.get_rank() |
| 128 | |
| 129 | device = torch.device(f"cuda:{ep_rank}") |
| 130 | |
| 131 | fml_layers = [ |
| 132 | make_fused_moe_layer(ep_rank, layer_idx, test_config).to(device) |
| 133 | for layer_idx in range(test_config.num_layers) |
| 134 | ] |
| 135 | rank_expert_weights = [fml.get_expert_weights() for fml in fml_layers] |
| 136 | |
| 137 | hidden_states = [] |
| 138 | router_logits = [] |
| 139 | for layer_idx in range(test_config.num_layers): |
| 140 | hidden_states.append( |
| 141 | torch.randn( |
| 142 | (test_config.num_tokens, test_config.hidden_size), |
| 143 | dtype=torch.bfloat16, |
| 144 | device=device, |
| 145 | ) |
| 146 | ) |
| 147 | router_logits.append( |
| 148 | torch.randn( |
| 149 | (test_config.num_tokens, test_config.num_experts), |
| 150 | dtype=torch.bfloat16, |
| 151 | device=device, |
| 152 | ) |
| 153 | ) |
| 154 | |
| 155 | out_before_shuffle = [] |
| 156 | with set_forward_context( |
| 157 | {}, |
| 158 | num_tokens=test_config.num_tokens, |
| 159 | num_tokens_across_dp=torch.tensor( |
| 160 | [test_config.num_tokens] * world_size, device="cpu", dtype=torch.int |
| 161 | ), |
| 162 | vllm_config=vllm_config, |
| 163 | ): |
| 164 | for lidx, fml in enumerate(fml_layers): |
| 165 | out_before_shuffle.append( |
| 166 | fml(hidden_states[lidx].clone(), router_logits[lidx].clone()) |
| 167 | ) |
| 168 | |
| 169 | indices = torch.zeros( |
| 170 | test_config.num_layers, test_config.num_experts, dtype=torch.long |
| 171 | ) |
| 172 | for lidx in range(test_config.num_layers): |
| 173 | indices[lidx] = torch.Tensor(range(test_config.num_experts)) |
| 174 | |
| 175 | shuffled_indices = torch.zeros_like(indices) |
| 176 | for lidx in range(test_config.num_layers): |
| 177 | shuffled_indices[lidx] = torch.randperm(test_config.num_experts) |
| 178 | |
| 179 | expert_buffer = [torch.empty_like(w) for w in rank_expert_weights[0]] |
| 180 | communicator = create_eplb_communicator( |
| 181 | group_coordinator=get_eplb_group(), |
| 182 | backend="torch_nccl", |
| 183 | expert_weights=rank_expert_weights, |
| 184 | expert_buffer=expert_buffer, |
| 185 | ) |
| 186 | rearrange_expert_weights_inplace( |
| 187 | indices, |
| 188 | shuffled_indices, |
| 189 | rank_expert_weights, |
| 190 | expert_buffer, |
| 191 | ep_group, |
| 192 | communicator, |
| 193 | ) |
| 194 | |
| 195 | num_global_experts = test_config.num_experts |
| 196 | |
| 197 | logical_to_physical_map_list = [] |
| 198 | for lidx, fml in enumerate(fml_layers): |
| 199 | physical_to_logical_map = shuffled_indices[lidx].to(device) |
| 200 | logical_to_physical_map = torch.empty( |
| 201 | (num_global_experts,), dtype=torch.int32, device=device |
| 202 | ) |
| 203 | logical_to_physical_map[physical_to_logical_map] = torch.arange( |
| 204 | 0, num_global_experts, dtype=torch.int32, device=device |
| 205 | ) |
| 206 | logical_to_physical_map_list.append( |
| 207 | logical_to_physical_map.reshape(num_global_experts, 1) |
| 208 | ) |
| 209 | |
| 210 | logical_to_physical_map = torch.stack(logical_to_physical_map_list) |
| 211 | |
| 212 | for lidx, fml in enumerate(fml_layers): |
| 213 | logical_replica_count = torch.ones( |
| 214 | (test_config.num_layers, num_global_experts), |
| 215 | dtype=torch.int32, |
| 216 | device=device, |
| 217 | ) |
| 218 | fml.eplb_state = EplbLayerState() |
| 219 | fml.set_eplb_state( |
| 220 | lidx, |
| 221 | torch.zeros( |
| 222 | (test_config.num_layers, num_global_experts), |
| 223 | dtype=torch.int32, |
| 224 | device=device, |
| 225 | ), |
| 226 | logical_to_physical_map, |
| 227 | logical_replica_count, |
| 228 | ) |
| 229 | fml.router.eplb_state.should_record_tensor = torch.ones( |
| 230 | (), dtype=torch.bool, device=device |
| 231 | ) |
| 232 | fml.router.eplb_state.num_unpadded_tokens_tensors = [ |
| 233 | torch.tensor(0, dtype=torch.int32, device=device) |
| 234 | ] |
| 235 | |
| 236 | out_after_shuffle = [] |
| 237 | with set_forward_context( |
| 238 | {}, |
| 239 | num_tokens=test_config.num_tokens, |
| 240 | num_tokens_across_dp=torch.tensor( |
| 241 | [test_config.num_tokens] * world_size, device="cpu", dtype=torch.int |
| 242 | ), |
| 243 | vllm_config=vllm_config, |
| 244 | ): |
| 245 | for lidx, fml in enumerate(fml_layers): |
| 246 | out_after_shuffle.append( |
| 247 | fml(hidden_states[lidx].clone(), router_logits[lidx].clone()) |
| 248 | ) |
| 249 | |
| 250 | for lidx in range(test_config.num_layers): |
| 251 | torch.testing.assert_close( |
| 252 | out_before_shuffle[lidx], out_after_shuffle[lidx], atol=1e-1, rtol=1e-1 |
| 253 | ) |
| 254 | |
| 255 | |
| 256 | @pytest.mark.parametrize("world_size", [2, 4]) |
| 257 | @pytest.mark.parametrize("num_layers", [8]) |
| 258 | @pytest.mark.parametrize("num_experts", [32]) |
| 259 | @pytest.mark.parametrize("hidden_size", [256]) |
| 260 | @pytest.mark.parametrize("intermediate_size", [256]) |
| 261 | @pytest.mark.parametrize("num_tokens", [256]) |
| 262 | @pytest.mark.parametrize("moe_backend", ["flashinfer_trtllm", "flashinfer_cutlass"]) |
| 263 | def test_eplb_fml( |
| 264 | world_size: int, |
| 265 | num_layers: int, |
| 266 | num_experts: int, |
| 267 | hidden_size: int, |
| 268 | intermediate_size: int, |
| 269 | num_tokens: int, |
| 270 | moe_backend: str, |
| 271 | ): |
| 272 | if torch.accelerator.device_count() < world_size: |
| 273 | pytest.skip(f"Need at least {world_size} GPUs to run the test") |
| 274 | if current_platform.is_xpu(): |
| 275 | pytest.skip("NVFP4 quantization and flashinfer backends are CUDA-only") |
| 276 | |
| 277 | num_local_experts = num_experts // world_size |
| 278 | num_topk = 4 |
| 279 | |
| 280 | test_config = TestConfig( |
| 281 | num_layers=num_layers, |
| 282 | num_experts=num_experts, |
| 283 | num_local_experts=num_local_experts, |
| 284 | num_topk=num_topk, |
| 285 | hidden_size=hidden_size, |
| 286 | intermediate_size=intermediate_size, |
| 287 | num_tokens=num_tokens, |
| 288 | moe_backend=moe_backend, |
| 289 | ) |
| 290 | |
| 291 | distributed_run( |
| 292 | _test_eplb_fml, |
| 293 | world_size, |
| 294 | test_config, |
| 295 | ) |