| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | import torch.nn as nn |
| 6 | from transformers import PreTrainedConfig |
| 7 | |
| 8 | from vllm import envs |
| 9 | from vllm.config.lora import LoRAConfig |
| 10 | from vllm.distributed.utils import divide |
| 11 | from vllm.lora.layers.base import BaseLayerWithLoRA |
| 12 | from vllm.model_executor.custom_op import maybe_get_oot_by_class |
| 13 | from vllm.model_executor.layers.fused_moe import MoERunner |
| 14 | from vllm.model_executor.layers.fused_moe.experts.lora_context import MoELoRAContext |
| 15 | from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import ( |
| 16 | LoRAExpertsMixin, |
| 17 | ) |
| 18 | from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import ( |
| 19 | FusedMoEModularMethod, |
| 20 | ) |
| 21 | from vllm.model_executor.layers.fused_moe.modular_kernel import ( |
| 22 | FusedMoEKernelModularImpl, |
| 23 | ) |
| 24 | from vllm.platforms import current_platform |
| 25 | |
| 26 | from .utils import _get_lora_aux_cuda_stream, _get_lora_device |
| 27 | |
| 28 | |
| 29 | class FusedMoEWithLoRA(BaseLayerWithLoRA): |
| 30 | def __init__(self, base_layer: MoERunner) -> None: |
| 31 | super().__init__() |
| 32 | self.base_layer = base_layer |
| 33 | self.moe_config = base_layer.moe_config |
| 34 | self._shared_experts = base_layer._shared_experts |
| 35 | self._ep_check() |
| 36 | |
| 37 | routed_experts = self.base_layer.routed_experts |
| 38 | routed_experts.lora_base_layer_prefix = "base_layer." |
| 39 | assert not routed_experts.quant_method.is_monolithic, ( |
| 40 | "Monolithic kernels are not supported for Fused MoE LoRA." |
| 41 | ) |
| 42 | |
| 43 | # Use the MoE-aware TP rank/size: when EP is active, MoERunner collapses |
| 44 | # moe_parallel_config.tp_size to 1 (experts are sharded across the |
| 45 | # TP group instead). |
| 46 | moe_parallel_config = self.moe_config.moe_parallel_config |
| 47 | self.tp_size = moe_parallel_config.tp_size |
| 48 | self.tp_rank = moe_parallel_config.tp_rank |
| 49 | self.device = _get_lora_device(base_layer) |
| 50 | |
| 51 | self._enable_aux_cuda_stream = envs.VLLM_LORA_ENABLE_DUAL_STREAM |
| 52 | self._init_lora_stream_context() |
| 53 | # For non-gated MoE (is_act_and_mul=False), only 1 slice is needed |
| 54 | # since there's only up_proj (w1), not gate_proj + up_proj (w1 + w3) |
| 55 | self._w13_slices = 2 if base_layer.moe_config.is_act_and_mul else 1 |
| 56 | # Set from lora_config.enable_moe_shared_loras in create_lora_weights. |
| 57 | # When True, w13 lora_A and w2 lora_B are stored once for all experts. |
| 58 | self.enable_moe_shared_loras = False |
| 59 | # Mirrors per-(lora_id) layout of `self.lora_a_stacked` (built in |
| 60 | # `create_lora_weights`) so `create_dummy_lora`'s n_slices fallback |
| 61 | # matches `lora_a_stacked` length under EP. |
| 62 | self.n_slices = self.local_num_experts * (self._w13_slices + 1) |
| 63 | |
| 64 | routed_experts._ensure_moe_quant_config_init() |
| 65 | if getattr(routed_experts.quant_method, "supports_internal_mk", False): |
| 66 | moe_kernel = routed_experts.quant_method.moe_kernel |
| 67 | assert moe_kernel is not None, ( |
| 68 | "Fused MoE quant method must provide a moe_kernel." |
| 69 | ) |
| 70 | # Don't let the kernel own shared experts so the runner can |
| 71 | # overlap them with routed experts via a separate CUDA stream. |
| 72 | assert isinstance(moe_kernel.impl, FusedMoEKernelModularImpl) |
| 73 | moe_kernel.impl.shared_experts = None |
| 74 | else: |
| 75 | raise RuntimeError("Old-style modular kernel not supported.") |
| 76 | |
| 77 | assert moe_kernel.supports_lora(), ( |
| 78 | f"{type(moe_kernel.fused_experts).__name__} does not support LoRA. " |
| 79 | "For unquantized MoE, set moe_backend='triton' or moe_backend='auto' " |
| 80 | "(auto selects Triton automatically when LoRA is enabled). " |
| 81 | "For quantized MoE, mix LoRAExpertsMixin into the experts class " |
| 82 | "and consume self._lora_context in apply()." |
| 83 | ) |
| 84 | self._moe_kernel = moe_kernel |
| 85 | self.base_layer._replace_quant_method( |
| 86 | FusedMoEModularMethod(self.base_layer._quant_method, moe_kernel) |
| 87 | ) |
| 88 | |
| 89 | @property |
| 90 | def hidden_size(self) -> int: |
| 91 | return self.moe_config.hidden_dim |
| 92 | |
| 93 | @property |
| 94 | def local_num_experts(self) -> int: |
| 95 | return self.moe_config.num_local_experts |
| 96 | |
| 97 | @property |
| 98 | def global_num_experts(self) -> int: |
| 99 | return self.moe_config.num_experts |
| 100 | |
| 101 | @property |
| 102 | def ep_rank(self) -> int: |
| 103 | return self.moe_config.moe_parallel_config.ep_rank |
| 104 | |
| 105 | @property |
| 106 | def use_ep(self) -> bool: |
| 107 | return self.moe_config.moe_parallel_config.use_ep |
| 108 | |
| 109 | @property |
| 110 | def intermediate_size_per_partition(self) -> int: |
| 111 | return self.moe_config.intermediate_size_per_partition |
| 112 | |
| 113 | def _init_lora_stream_context(self) -> None: |
| 114 | self._lora_stream: torch.cuda.Stream | None = None |
| 115 | self._events: tuple[torch.cuda.Event, ...] | None = None |
| 116 | if not self._enable_aux_cuda_stream: |
| 117 | return |
| 118 | if not current_platform.is_cuda_alike(): |
| 119 | return |
| 120 | self._lora_stream = _get_lora_aux_cuda_stream() |
| 121 | # 4 events: 2 per (base GEMM, LoRA) pair so w13 and w2 don't reuse |
| 122 | # the same event objects; reuse-within-a-pair is fine because the |
| 123 | # second pair starts only after intermediate_cache1.add_() has joined. |
| 124 | self._events = tuple(torch.cuda.Event() for _ in range(4)) |
| 125 | |
| 126 | def _build_lora_context(self): |
| 127 | use_dual_stream = ( |
| 128 | self._enable_aux_cuda_stream |
| 129 | and not self.fully_sharded |
| 130 | and self._lora_stream is not None |
| 131 | ) |
| 132 | return MoELoRAContext( |
| 133 | w13_lora_a_stacked=self.w13_lora_a_stacked, |
| 134 | w13_lora_b_stacked=self.w13_lora_b_stacked, |
| 135 | w2_lora_a_stacked=self.w2_lora_a_stacked, |
| 136 | w2_lora_b_stacked=self.w2_lora_b_stacked, |
| 137 | adapter_enabled=self.adapter_enabled, |
| 138 | max_loras=self.max_loras, |
| 139 | top_k=self.moe_config.experts_per_token, |
| 140 | w13_num_slices=self._w13_slices, |
| 141 | fully_sharded=self.fully_sharded, |
| 142 | tp_rank=self.tp_rank, |
| 143 | tp_size=self.tp_size, |
| 144 | local_num_experts=self.local_num_experts, |
| 145 | punica_wrapper=self.punica_wrapper, |
| 146 | use_tuned_config=bool(envs.VLLM_TUNED_CONFIG_FOLDER), |
| 147 | enable_moe_shared_loras=self.enable_moe_shared_loras, |
| 148 | aux_stream=self._lora_stream if use_dual_stream else None, |
| 149 | events=self._events if use_dual_stream else None, |
| 150 | ) |
| 151 | |
| 152 | @property |
| 153 | def _w13_a_num_experts(self) -> int: |
| 154 | """Expert-dim of the w13 lora_A buffer: 1 when shared.""" |
| 155 | return 1 if self.enable_moe_shared_loras else self.local_num_experts |
| 156 | |
| 157 | @property |
| 158 | def _w2_b_num_experts(self) -> int: |
| 159 | """Expert-dim of the w2 lora_B buffer: 1 when shared.""" |
| 160 | return 1 if self.enable_moe_shared_loras else self.local_num_experts |
| 161 | |
| 162 | def _create_lora_a_weights( |
| 163 | self, |
| 164 | max_loras: int, |
| 165 | lora_config: LoRAConfig, |
| 166 | ): |
| 167 | self.w13_lora_a_stacked: tuple[torch.Tensor, ...] = tuple( |
| 168 | torch.zeros( |
| 169 | ( |
| 170 | max_loras, |
| 171 | self._w13_a_num_experts, |
| 172 | lora_config.max_lora_rank |
| 173 | if not self.fully_sharded |
| 174 | else divide(lora_config.max_lora_rank, self.tp_size), |
| 175 | self.hidden_size, |
| 176 | ), |
| 177 | dtype=lora_config.lora_dtype, |
| 178 | device=self.device, |
| 179 | ) |
| 180 | for _ in range(self._w13_slices) |
| 181 | ) |
| 182 | self.w2_lora_a_stacked: tuple[torch.Tensor, ...] = ( |
| 183 | torch.zeros( |
| 184 | ( |
| 185 | max_loras, |
| 186 | self.local_num_experts, |
| 187 | lora_config.max_lora_rank, |
| 188 | self.intermediate_size_per_partition, |
| 189 | ), |
| 190 | dtype=lora_config.lora_dtype, |
| 191 | device=self.device, |
| 192 | ), |
| 193 | ) |
| 194 | |
| 195 | def _create_lora_b_weights(self, max_loras: int, lora_config: LoRAConfig): |
| 196 | self.w13_lora_b_stacked: tuple[torch.Tensor, ...] = tuple( |
| 197 | torch.zeros( |
| 198 | ( |
| 199 | max_loras, |
| 200 | self.local_num_experts, |
| 201 | self.intermediate_size_per_partition, |
| 202 | lora_config.max_lora_rank, |
| 203 | ), |
| 204 | dtype=lora_config.lora_dtype, |
| 205 | device=self.device, |
| 206 | ) |
| 207 | for _ in range(self._w13_slices) |
| 208 | ) |
| 209 | self.w2_lora_b_stacked: tuple[torch.Tensor, ...] = ( |
| 210 | torch.zeros( |
| 211 | ( |
| 212 | max_loras, |
| 213 | self._w2_b_num_experts, |
| 214 | self.hidden_size |
| 215 | if not self.fully_sharded |
| 216 | else divide(self.hidden_size, self.tp_size), |
| 217 | lora_config.max_lora_rank, |
| 218 | ), |
| 219 | dtype=lora_config.lora_dtype, |
| 220 | device=self.device, |
| 221 | ), |
| 222 | ) |
| 223 | |
| 224 | def _ep_check(self): |
| 225 | if self.use_ep: |
| 226 | moe_config = self.moe_config |
| 227 | all2all_backend = moe_config.moe_parallel_config.all2all_backend |
| 228 | assert all2all_backend == "allgather_reducescatter", ( |
| 229 | "Fused MoE LoRA with EP currently only supports " |
| 230 | f"all2all_backend='allgather_reducescatter', got '{all2all_backend}'." |
| 231 | ) |
| 232 | assert not moe_config.moe_parallel_config.is_sequence_parallel |
| 233 | |
| 234 | def _verify_ep_fs(self, lora_config: LoRAConfig): |
| 235 | # EP and fully_sharded LoRA both partition along the same TP group — |
| 236 | # EP on the expert dim, fully_sharded on the LoRA rank dim — with |
| 237 | # mutually contradictory assumptions about which rank holds which |
| 238 | # expert's rank-shard. |
| 239 | assert not (self.use_ep and lora_config.fully_sharded_loras), ( |
| 240 | "Fused MoE LoRA does not support enable_expert_parallel=True " |
| 241 | "together with fully_sharded_loras=True. Disable one of them." |
| 242 | ) |
| 243 | |
| 244 | def create_lora_weights( |
| 245 | self, |
| 246 | max_loras: int, |
| 247 | lora_config: LoRAConfig, |
| 248 | model_config: PreTrainedConfig | None = None, |
| 249 | ) -> None: |
| 250 | """Initializes lora matrices.""" |
| 251 | self._verify_ep_fs(lora_config) |
| 252 | self.max_loras = lora_config.max_loras |
| 253 | self.fully_sharded = lora_config.fully_sharded_loras |
| 254 | self.enable_moe_shared_loras = lora_config.enable_moe_shared_loras |
| 255 | |
| 256 | self.adapter_enabled = torch.tensor( |
| 257 | [0] * (max_loras + 1), dtype=torch.int, device=self.device |
| 258 | ) |
| 259 | |
| 260 | self._create_lora_a_weights(max_loras, lora_config) |
| 261 | self._create_lora_b_weights(max_loras, lora_config) |
| 262 | # They will be used by 'LoRALayerWeights.create_dummy_lora_weights' |
| 263 | # to create a dummy LoRA weights. |
| 264 | # TODO Optimize this section |
| 265 | self.lora_a_stacked = [] |
| 266 | self.lora_b_stacked = [] |
| 267 | for lora_id in range(max_loras): |
| 268 | for experts_id in range(self.local_num_experts): |
| 269 | # For gated MoE: gate_proj (w1), down_proj (w2), up_proj (w3) |
| 270 | # For non-gated MoE: up_proj (w1), down_proj (w2) |
| 271 | # Shared factors are collapsed (expert-dim 1): index 0. |
| 272 | w13_a_eid = 0 if self.enable_moe_shared_loras else experts_id |
| 273 | w2_b_eid = 0 if self.enable_moe_shared_loras else experts_id |
| 274 | self.lora_a_stacked.append( |
| 275 | self.w13_lora_a_stacked[0][lora_id][w13_a_eid] |
| 276 | ) |
| 277 | self.lora_a_stacked.append( |
| 278 | self.w2_lora_a_stacked[0][lora_id][experts_id] |
| 279 | ) |
| 280 | |
| 281 | self.lora_b_stacked.append( |
| 282 | self.w13_lora_b_stacked[0][lora_id][experts_id] |
| 283 | ) |
| 284 | self.lora_b_stacked.append(self.w2_lora_b_stacked[0][lora_id][w2_b_eid]) |
| 285 | |
| 286 | # Only add w3 (up_proj) for gated MoE (_w13_slices == 2) |
| 287 | if self._w13_slices == 2: |
| 288 | self.lora_a_stacked.append( |
| 289 | self.w13_lora_a_stacked[1][lora_id][w13_a_eid] |
| 290 | ) |
| 291 | self.lora_b_stacked.append( |
| 292 | self.w13_lora_b_stacked[1][lora_id][experts_id] |
| 293 | ) |
| 294 | |
| 295 | def _slice_w13_a(self, w13_lora_a: torch.Tensor) -> torch.Tensor: |
| 296 | """Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA.""" |
| 297 | if self.tp_size == 1 or not self.fully_sharded: |
| 298 | return w13_lora_a |
| 299 | |
| 300 | # w13_lora_a shape (num_experts,rank,input_size) |
| 301 | current_lora_rank = w13_lora_a.shape[1] |
| 302 | assert current_lora_rank % self.tp_size == 0 |
| 303 | # Based on S-LoRA, we slice W13/W1/W3 A along the rank dim. |
| 304 | shard_size = self.w13_lora_a_stacked[0].shape[2] |
| 305 | start_idx = self.tp_rank * shard_size |
| 306 | end_idx = (self.tp_rank + 1) * shard_size |
| 307 | return w13_lora_a[:, start_idx:end_idx, :] |
| 308 | |
| 309 | def _slice_w13_b(self, w13_lora_b: torch.Tensor): |
| 310 | if self.tp_size == 1: |
| 311 | return w13_lora_b |
| 312 | |
| 313 | # w13_lora_b shape (num_experts,output_size,rank) |
| 314 | shard_size = self.intermediate_size_per_partition |
| 315 | start_idx = self.tp_rank * shard_size |
| 316 | end_idx = (self.tp_rank + 1) * shard_size |
| 317 | |
| 318 | return w13_lora_b[:, start_idx:end_idx, :] |
| 319 | |
| 320 | def _slice_w2_a(self, w2_lora_a: torch.Tensor) -> torch.Tensor: |
| 321 | """Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA.""" |
| 322 | if self.tp_size == 1: |
| 323 | return w2_lora_a |
| 324 | # w2_lora_a shape (num_experts,rank,input_size) |
| 325 | shard_size = self.intermediate_size_per_partition |
| 326 | start_idx = self.tp_rank * shard_size |
| 327 | end_idx = (self.tp_rank + 1) * shard_size |
| 328 | |
| 329 | return w2_lora_a[:, :, start_idx:end_idx] |
| 330 | |
| 331 | def _slice_w2_b(self, w2_lora_b: torch.Tensor) -> torch.Tensor: |
| 332 | """Applies to FusedMoEWithLoRA and FusedMoE3DWithLoRA.""" |
| 333 | if self.tp_size == 1 or not self.fully_sharded: |
| 334 | return w2_lora_b |
| 335 | # Based on S-LoRA, we slice W2 B along the hidden_size dim. |
| 336 | # w2_lora_b shape (num_experts,output_size,rank) |
| 337 | shard_size = self.w2_lora_b_stacked[0].shape[2] |
| 338 | start_idx = self.tp_rank * shard_size |
| 339 | end_idx = (self.tp_rank + 1) * shard_size |
| 340 | |
| 341 | return w2_lora_b[:, start_idx:end_idx, :] |
| 342 | |
| 343 | def _match_expert_dim( |
| 344 | self, src: torch.Tensor, buffer: torch.Tensor |
| 345 | ) -> torch.Tensor: |
| 346 | """Align a LoRA factor's expert-dim (dim 0) to its stacked buffer. |
| 347 | |
| 348 | Equal dims pass through. When the buffer is collapsed to expert-dim 1 |
| 349 | (a shared factor) but the source carries per-expert copies, keep |
| 350 | the first — every expert shares the same factor, so the copies are |
| 351 | identical for a real adapter and irrelevant (zeros) for the dummy. |
| 352 | """ |
| 353 | tgt = buffer.shape[1] |
| 354 | if src.shape[0] == tgt: |
| 355 | return src |
| 356 | assert tgt == 1, f"expert-dim mismatch: source {src.shape[0]} vs buffer {tgt}" |
| 357 | return src[:1] |
| 358 | |
| 359 | def reset_lora(self, index: int): |
| 360 | """Resets the lora weights at index back to 0.""" |
| 361 | for pos in range(self._w13_slices): |
| 362 | self.w13_lora_a_stacked[pos][index] = 0 |
| 363 | self.w13_lora_b_stacked[pos][index] = 0 |
| 364 | |
| 365 | self.w2_lora_a_stacked[0][index] = 0 |
| 366 | self.w2_lora_b_stacked[0][index] = 0 |
| 367 | self.adapter_enabled[index] = 0 |
| 368 | |
| 369 | # |
| 370 | |
| 371 | def set_lora( |
| 372 | self, |
| 373 | index: int, |
| 374 | lora_a: torch.Tensor | list[torch.Tensor], |
| 375 | lora_b: torch.Tensor | list[torch.Tensor], |
| 376 | ): |
| 377 | """Overwrites lora tensors at index.""" |
| 378 | # Make mypy happy |
| 379 | assert isinstance(lora_a, list) |
| 380 | assert isinstance(lora_b, list) |
| 381 | |
| 382 | self.reset_lora(index) |
| 383 | self.adapter_enabled[index] = 1 |
| 384 | |
| 385 | w1_lora_a, w2_lora_a, w3_lora_a = lora_a |
| 386 | w1_lora_b, w2_lora_b, w3_lora_b = lora_b |
| 387 | |
| 388 | # EP slicing is done once at add time in |
| 389 | # LoRAModelManager._slice_moe_lora_ep, so by here the cached |
| 390 | # tensors already match the local-expert dim of the stacked buffers. |
| 391 | # For shared-loras adapters the w13 lora_A and w2 lora_B buffers are |
| 392 | # collapsed to expert-dim 1: match each factor to the expert-dim of |
| 393 | # the buffer it is copied into. A real shared factor already has |
| 394 | # expert-dim 1 (no-op); a per-expert source (e.g. the warmup dummy |
| 395 | # LoRA, which always stacks to num_experts) is collapsed to the shared |
| 396 | # slot since every expert shares the same factor. |
| 397 | w1_lora_a = self._match_expert_dim(w1_lora_a, self.w13_lora_a_stacked[0]) |
| 398 | w3_lora_a = self._match_expert_dim(w3_lora_a, self.w13_lora_a_stacked[-1]) |
| 399 | w2_lora_a = self._match_expert_dim(w2_lora_a, self.w2_lora_a_stacked[0]) |
| 400 | w1_lora_b = self._match_expert_dim(w1_lora_b, self.w13_lora_b_stacked[0]) |
| 401 | w3_lora_b = self._match_expert_dim(w3_lora_b, self.w13_lora_b_stacked[-1]) |
| 402 | w2_lora_b = self._match_expert_dim(w2_lora_b, self.w2_lora_b_stacked[0]) |
| 403 | |
| 404 | slliced_w1_lora_a = self._slice_w13_a(w1_lora_a) |
| 405 | slliced_w1_lora_b = self._slice_w13_b(w1_lora_b) |
| 406 | |
| 407 | sliced_w2_lora_a = self._slice_w2_a(w2_lora_a) |
| 408 | sliced_w2_lora_b = self._slice_w2_b(w2_lora_b) |
| 409 | |
| 410 | self.w13_lora_a_stacked[0][ |
| 411 | index, :, : slliced_w1_lora_a.shape[1], : slliced_w1_lora_a.shape[2] |
| 412 | ].copy_(slliced_w1_lora_a, non_blocking=True) |
| 413 | |
| 414 | self.w13_lora_b_stacked[0][ |
| 415 | index, :, : slliced_w1_lora_b.shape[1], : slliced_w1_lora_b.shape[2] |
| 416 | ].copy_(slliced_w1_lora_b, non_blocking=True) |
| 417 | |
| 418 | # Only copy w3 (up_proj) for gated MoE (_w13_slices == 2) |
| 419 | if self._w13_slices == 2: |
| 420 | slliced_w3_lora_a = self._slice_w13_a(w3_lora_a) |
| 421 | slliced_w3_lora_b = self._slice_w13_b(w3_lora_b) |
| 422 | |
| 423 | self.w13_lora_a_stacked[1][ |
| 424 | index, :, : slliced_w3_lora_a.shape[1], : slliced_w3_lora_a.shape[2] |
| 425 | ].copy_(slliced_w3_lora_a, non_blocking=True) |
| 426 | |
| 427 | self.w13_lora_b_stacked[1][ |
| 428 | index, :, : slliced_w3_lora_b.shape[1], : slliced_w3_lora_b.shape[2] |
| 429 | ].copy_(slliced_w3_lora_b, non_blocking=True) |
| 430 | |
| 431 | self.w2_lora_a_stacked[0][ |
| 432 | index, :, : sliced_w2_lora_a.shape[1], : sliced_w2_lora_a.shape[2] |
| 433 | ].copy_(sliced_w2_lora_a, non_blocking=True) |
| 434 | |
| 435 | self.w2_lora_b_stacked[0][ |
| 436 | index, :, : sliced_w2_lora_b.shape[1], : sliced_w2_lora_b.shape[2] |
| 437 | ].copy_(sliced_w2_lora_b, non_blocking=True) |
| 438 | |
| 439 | def set_mapping(self, punica_wrapper): |
| 440 | super().set_mapping(punica_wrapper) |
| 441 | lora_context = self._build_lora_context() |
| 442 | fused_experts = self._moe_kernel.fused_experts |
| 443 | assert isinstance(fused_experts, LoRAExpertsMixin), ( |
| 444 | f"{type(fused_experts).__name__} does not support LoRA context setup." |
| 445 | ) |
| 446 | fused_experts.set_lora_context(lora_context) |
| 447 | prepare_finalize = self._moe_kernel.prepare_finalize |
| 448 | if hasattr(prepare_finalize, "set_lora_context"): |
| 449 | prepare_finalize.set_lora_context(lora_context) |
| 450 | |
| 451 | def forward(self, *args, **kwargs): |
| 452 | return self.base_layer.forward(*args, **kwargs) |
| 453 | |
| 454 | @property |
| 455 | def quant_method(self): |
| 456 | return self.base_layer._quant_method |
| 457 | |
| 458 | @property |
| 459 | def runner(self) -> MoERunner: |
| 460 | return self.base_layer |
| 461 | |
| 462 | @classmethod |
| 463 | def can_replace_layer( |
| 464 | cls, |
| 465 | source_layer: nn.Module, |
| 466 | lora_config: LoRAConfig, |
| 467 | packed_modules_list: list, |
| 468 | model_config: PreTrainedConfig | None = None, |
| 469 | ) -> bool: |
| 470 | """Returns True if the layer can be replaced by this LoRA layer.""" |
| 471 | # source_layer is MoERunner |
| 472 | moe_cls = maybe_get_oot_by_class(MoERunner) |
| 473 | return isinstance(source_layer, moe_cls) and len(packed_modules_list) == 2 |
| 474 | |
| 475 | |
| 476 | class FusedMoE3DWithLoRA(FusedMoEWithLoRA): |
| 477 | def __init__(self, base_layer: MoERunner): |
| 478 | super().__init__(base_layer) |
| 479 | self._w13_slices = 1 |
| 480 | |
| 481 | def _create_lora_b_weights(self, max_loras, lora_config): |
| 482 | self.w13_lora_b_stacked: tuple[torch.Tensor] = tuple( |
| 483 | torch.zeros( |
| 484 | ( |
| 485 | max_loras, |
| 486 | self.local_num_experts, |
| 487 | self.intermediate_size_per_partition * 2, |
| 488 | lora_config.max_lora_rank, |
| 489 | ), |
| 490 | dtype=lora_config.lora_dtype, |
| 491 | device=self.device, |
| 492 | ) |
| 493 | for _ in range(self._w13_slices) |
| 494 | ) |
| 495 | self.w2_lora_b_stacked: tuple[torch.Tensor] = ( |
| 496 | torch.zeros( |
| 497 | ( |
| 498 | max_loras, |
| 499 | self.local_num_experts, |
| 500 | self.hidden_size |
| 501 | if not self.fully_sharded |
| 502 | else divide(self.hidden_size, self.tp_size), |
| 503 | lora_config.max_lora_rank, |
| 504 | ), |
| 505 | dtype=lora_config.lora_dtype, |
| 506 | device=self.device, |
| 507 | ), |
| 508 | ) |
| 509 | |
| 510 | def create_lora_weights( |
| 511 | self, |
| 512 | max_loras: int, |
| 513 | lora_config: LoRAConfig, |
| 514 | model_config: PreTrainedConfig | None = None, |
| 515 | ) -> None: |
| 516 | """Initializes lora matrices.""" |
| 517 | if model_config is None: |
| 518 | raise ValueError("model_config must be provided for MoE LoRA.") |
| 519 | architectures = model_config.architectures |
| 520 | if not architectures: |
| 521 | raise ValueError("model_config.architectures must be defined for MoE LoRA.") |
| 522 | self._verify_ep_fs(lora_config) |
| 523 | self._base_model = architectures[0] |
| 524 | self.max_loras = lora_config.max_loras |
| 525 | self.fully_sharded = lora_config.fully_sharded_loras |
| 526 | |
| 527 | self.adapter_enabled = torch.tensor( |
| 528 | [0] * (max_loras + 1), dtype=torch.int, device=self.device |
| 529 | ) |
| 530 | |
| 531 | self._create_lora_a_weights(max_loras, lora_config) |
| 532 | self._create_lora_b_weights(max_loras, lora_config) |
| 533 | |
| 534 | def _slice_w13_b(self, w13_lora_b: torch.Tensor): |
| 535 | if self.tp_size == 1: |
| 536 | return w13_lora_b |
| 537 | |
| 538 | # w13_lora_b shape (num_experts,output_size,rank) |
| 539 | shard_size = self.intermediate_size_per_partition |
| 540 | start_idx = self.tp_rank * shard_size |
| 541 | end_idx = (self.tp_rank + 1) * shard_size |
| 542 | # HACK: Currently, only GPT-OSS is in interleaved order |
| 543 | if self._base_model == "GptOssForCausalLM": |
| 544 | # For models like GPT-OSS, the weights of w1 (gate_proj) and w3 (up_proj) |
| 545 | # in the interleaved order, and corresponding LoRA need to be processed. |
| 546 | w1_lora_b = w13_lora_b[:, ::2, :] |
| 547 | w3_lora_b = w13_lora_b[:, 1::2, :] |
| 548 | sliced_w1_lora_b = w1_lora_b[:, start_idx:end_idx, :] |
| 549 | sliced_w3_lora_b = w3_lora_b[:, start_idx:end_idx, :] |
| 550 | |
| 551 | return torch.stack([sliced_w1_lora_b, sliced_w3_lora_b], dim=2).flatten( |
| 552 | 1, 2 |
| 553 | ) |
| 554 | else: |
| 555 | slice_size = w13_lora_b.shape[1] // 2 |
| 556 | w1_lora_b = w13_lora_b[:, :slice_size, :] |
| 557 | w3_lora_b = w13_lora_b[:, slice_size:, :] |
| 558 | sliced_w1_lora_b = w1_lora_b[:, start_idx:end_idx, :] |
| 559 | sliced_w3_lora_b = w3_lora_b[:, start_idx:end_idx, :] |
| 560 | |
| 561 | return torch.cat([sliced_w1_lora_b, sliced_w3_lora_b], dim=1) |
| 562 | |
| 563 | def set_lora( |
| 564 | self, |
| 565 | index: int, |
| 566 | lora_a: torch.Tensor | list[torch.Tensor], |
| 567 | lora_b: torch.Tensor | list[torch.Tensor], |
| 568 | ): |
| 569 | """Overwrites lora tensors at index.""" |
| 570 | # Make mypy happy |
| 571 | assert isinstance(lora_a, list) |
| 572 | assert isinstance(lora_b, list) |
| 573 | assert len(lora_a) == len(lora_b) == 2 |
| 574 | |
| 575 | self.reset_lora(index) |
| 576 | self.adapter_enabled[index] = 1 |
| 577 | |
| 578 | w13_lora_a, w2_lora_a = lora_a |
| 579 | w13_lora_b, w2_lora_b = lora_b |
| 580 | |
| 581 | sliced_w13_lora_a = self._slice_w13_a(w13_lora_a) |
| 582 | sliced_w13_lora_b = self._slice_w13_b(w13_lora_b) |
| 583 | |
| 584 | sliced_w2_lora_a = self._slice_w2_a(w2_lora_a) |
| 585 | sliced_w2_lora_b = self._slice_w2_b(w2_lora_b) |
| 586 | |
| 587 | self.w13_lora_a_stacked[0][ |
| 588 | index, :, : sliced_w13_lora_a.shape[1], : sliced_w13_lora_a.shape[2] |
| 589 | ].copy_(sliced_w13_lora_a, non_blocking=True) |
| 590 | self.w2_lora_a_stacked[0][ |
| 591 | index, :, : sliced_w2_lora_a.shape[1], : sliced_w2_lora_a.shape[2] |
| 592 | ].copy_(sliced_w2_lora_a, non_blocking=True) |
| 593 | |
| 594 | self.w13_lora_b_stacked[0][ |
| 595 | index, :, : sliced_w13_lora_b.shape[1], : sliced_w13_lora_b.shape[2] |
| 596 | ].copy_(sliced_w13_lora_b, non_blocking=True) |
| 597 | self.w2_lora_b_stacked[0][ |
| 598 | index, :, : sliced_w2_lora_b.shape[1], : sliced_w2_lora_b.shape[2] |
| 599 | ].copy_(sliced_w2_lora_b, non_blocking=True) |
| 600 | |
| 601 | @property |
| 602 | def w13_input_size(self): |
| 603 | """Full size.""" |
| 604 | return self.w13_lora_a_stacked[0].shape[-1] |
| 605 | |
| 606 | @property |
| 607 | def w13_output_size(self): |
| 608 | """Full size.""" |
| 609 | return self.w13_lora_b_stacked[0].shape[-2] * self.tp_size |
| 610 | |
| 611 | @property |
| 612 | def w2_input_size(self): |
| 613 | """Full size.""" |
| 614 | return self.w2_lora_a_stacked[0].shape[-1] * self.tp_size |
| 615 | |
| 616 | @property |
| 617 | def w2_output_size(self): |
| 618 | """Full size.""" |
| 619 | return self.hidden_size |
| 620 | |
| 621 | @classmethod |
| 622 | def can_replace_layer( |
| 623 | cls, |
| 624 | source_layer: nn.Module, |
| 625 | lora_config: LoRAConfig, |
| 626 | packed_modules_list: list, |
| 627 | model_config: PreTrainedConfig | None = None, |
| 628 | ) -> bool: |
| 629 | """Returns True if the layer can be replaced by this LoRA layer.""" |
| 630 | # source_layer is MoERunner |
| 631 | moe_cls = maybe_get_oot_by_class(MoERunner) |
| 632 | return isinstance(source_layer, moe_cls) and len(packed_modules_list) == 1 |