| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from abc import abstractmethod |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.model_executor.layers.fused_moe import ( |
| 9 | FusedMoEMethodBase, |
| 10 | RoutedExperts, |
| 11 | SharedExperts, |
| 12 | ) |
| 13 | from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput |
| 14 | from vllm.model_executor.model_loader.reload.layerwise import ( |
| 15 | initialize_online_processing, |
| 16 | ) |
| 17 | from vllm.model_executor.utils import set_weight_attrs |
| 18 | |
| 19 | |
| 20 | class OnlineMoEMethodBase(FusedMoEMethodBase): |
| 21 | """Base for MoE methods that load full-precision weights on meta device |
| 22 | and quantize them after loading via the QeRL layerwise processing system. |
| 23 | """ |
| 24 | |
| 25 | uses_meta_device: bool = True |
| 26 | |
| 27 | def create_weights( |
| 28 | self, |
| 29 | layer: torch.nn.Module, |
| 30 | num_experts: int, |
| 31 | hidden_size: int, |
| 32 | intermediate_size_per_partition: int, |
| 33 | params_dtype: torch.dtype, |
| 34 | **extra_weight_attrs, |
| 35 | ): |
| 36 | layer.num_experts = num_experts |
| 37 | layer.orig_dtype = params_dtype |
| 38 | layer.weight_block_size = None |
| 39 | |
| 40 | # Fused gate_up_proj (column parallel) — full precision on meta device |
| 41 | w13_weight = torch.nn.Parameter( |
| 42 | torch.empty( |
| 43 | num_experts, |
| 44 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 45 | hidden_size, |
| 46 | device="meta", |
| 47 | dtype=params_dtype, |
| 48 | ), |
| 49 | requires_grad=False, |
| 50 | ) |
| 51 | layer.register_parameter("w13_weight", w13_weight) |
| 52 | set_weight_attrs(w13_weight, extra_weight_attrs) |
| 53 | |
| 54 | # down_proj (row parallel) — full precision on meta device |
| 55 | w2_weight = torch.nn.Parameter( |
| 56 | torch.empty( |
| 57 | num_experts, |
| 58 | hidden_size, |
| 59 | intermediate_size_per_partition, |
| 60 | device="meta", |
| 61 | dtype=params_dtype, |
| 62 | ), |
| 63 | requires_grad=False, |
| 64 | ) |
| 65 | layer.register_parameter("w2_weight", w2_weight) |
| 66 | set_weight_attrs(w2_weight, extra_weight_attrs) |
| 67 | |
| 68 | # BIASES (for models like GPT-OSS that have biased MoE) |
| 69 | if self.moe.has_bias: |
| 70 | w13_bias = torch.nn.Parameter( |
| 71 | torch.zeros( |
| 72 | num_experts, |
| 73 | self.moe.w13_num_shards * intermediate_size_per_partition, |
| 74 | device="meta", |
| 75 | dtype=layer.orig_dtype, |
| 76 | ), |
| 77 | requires_grad=False, |
| 78 | ) |
| 79 | layer.register_parameter("w13_bias", w13_bias) |
| 80 | set_weight_attrs(w13_bias, extra_weight_attrs) |
| 81 | |
| 82 | w2_bias = torch.nn.Parameter( |
| 83 | torch.zeros( |
| 84 | num_experts, |
| 85 | hidden_size, |
| 86 | device="meta", |
| 87 | dtype=layer.orig_dtype, |
| 88 | ), |
| 89 | requires_grad=False, |
| 90 | ) |
| 91 | layer.register_parameter("w2_bias", w2_bias) |
| 92 | set_weight_attrs(w2_bias, extra_weight_attrs) |
| 93 | |
| 94 | layer.w13_input_scale = None |
| 95 | layer.w2_input_scale = None |
| 96 | |
| 97 | initialize_online_processing(layer) |
| 98 | |
| 99 | def _zero_padding(self, layer: torch.nn.Module) -> None: |
| 100 | hidden_size = layer.moe_config.hidden_dim_unpadded |
| 101 | intermediate_size = layer.moe_config.intermediate_size_per_partition_unpadded |
| 102 | |
| 103 | w13_shard = layer.w13_weight.shape[1] // self.moe.w13_num_shards |
| 104 | if w13_shard > intermediate_size: |
| 105 | for shard in range(self.moe.w13_num_shards): |
| 106 | start = shard * w13_shard + intermediate_size |
| 107 | layer.w13_weight[:, start : (shard + 1) * w13_shard, :] = 0 |
| 108 | if layer.w13_weight.shape[2] > hidden_size: |
| 109 | layer.w13_weight[:, :, hidden_size:] = 0 |
| 110 | |
| 111 | if layer.w2_weight.shape[1] > hidden_size: |
| 112 | layer.w2_weight[:, hidden_size:, :] = 0 |
| 113 | if layer.w2_weight.shape[2] > intermediate_size: |
| 114 | layer.w2_weight[:, :, intermediate_size:] = 0 |
| 115 | |
| 116 | if getattr(layer, "w13_bias", None) is not None: |
| 117 | w13_bias_shard = layer.w13_bias.shape[1] // self.moe.w13_num_shards |
| 118 | if w13_bias_shard > intermediate_size: |
| 119 | for shard in range(self.moe.w13_num_shards): |
| 120 | start = shard * w13_bias_shard + intermediate_size |
| 121 | layer.w13_bias[:, start : (shard + 1) * w13_bias_shard] = 0 |
| 122 | |
| 123 | if ( |
| 124 | getattr(layer, "w2_bias", None) is not None |
| 125 | and layer.w2_bias.shape[1] > hidden_size |
| 126 | ): |
| 127 | layer.w2_bias[:, hidden_size:] = 0 |
| 128 | |
| 129 | @abstractmethod |
| 130 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 131 | pass |
| 132 | |
| 133 | @property |
| 134 | def supports_eplb(self) -> bool: |
| 135 | return True |
| 136 | |
| 137 | def apply_monolithic( |
| 138 | self, |
| 139 | layer: RoutedExperts, |
| 140 | x: torch.Tensor, |
| 141 | router_logits: torch.Tensor, |
| 142 | input_ids: torch.Tensor | None = None, |
| 143 | ) -> torch.Tensor | UnfinalizedMoEOutput: |
| 144 | assert self.is_monolithic |
| 145 | assert self.moe_kernel is not None |
| 146 | return self.moe_kernel.apply_monolithic( |
| 147 | x, |
| 148 | layer.w13_weight, |
| 149 | layer.w2_weight, |
| 150 | router_logits, |
| 151 | activation=layer.activation, |
| 152 | global_num_experts=layer.global_num_experts, |
| 153 | expert_map=layer.expert_map, |
| 154 | apply_router_weight_on_input=layer.apply_router_weight_on_input, |
| 155 | num_expert_group=layer.num_expert_group, |
| 156 | topk_group=layer.topk_group, |
| 157 | e_score_correction_bias=layer.e_score_correction_bias, |
| 158 | routed_scaling_factor=layer.routed_scaling_factor, |
| 159 | ) |
| 160 | |
| 161 | def apply( |
| 162 | self, |
| 163 | layer: RoutedExperts, |
| 164 | x: torch.Tensor, |
| 165 | topk_weights: torch.Tensor, |
| 166 | topk_ids: torch.Tensor, |
| 167 | shared_experts: SharedExperts | None, |
| 168 | shared_experts_input: torch.Tensor | None, |
| 169 | ) -> torch.Tensor | UnfinalizedMoEOutput: |
| 170 | assert not self.is_monolithic |
| 171 | assert self.moe_kernel is not None |
| 172 | return self.moe_kernel.apply( |
| 173 | x, |
| 174 | layer.w13_weight, |
| 175 | layer.w2_weight, |
| 176 | topk_weights, |
| 177 | topk_ids, |
| 178 | activation=layer.activation, |
| 179 | global_num_experts=layer.global_num_experts, |
| 180 | expert_map=layer.expert_map, |
| 181 | apply_router_weight_on_input=layer.apply_router_weight_on_input, |
| 182 | shared_experts=shared_experts, |
| 183 | shared_experts_input=shared_experts_input, |
| 184 | ) |