| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import typing |
| 4 | from collections.abc import Callable, Iterable |
| 5 | |
| 6 | import torch |
| 7 | import torch.distributed as dist |
| 8 | from torch import nn |
| 9 | from transformers import GptOssConfig |
| 10 | |
| 11 | from vllm.compilation.decorators import support_torch_compile |
| 12 | from vllm.config import CacheConfig, VllmConfig |
| 13 | from vllm.distributed import ( |
| 14 | get_dp_group, |
| 15 | get_ep_group, |
| 16 | get_pcp_group, |
| 17 | get_pp_group, |
| 18 | get_tensor_model_parallel_rank, |
| 19 | get_tensor_model_parallel_world_size, |
| 20 | tensor_model_parallel_all_gather, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.attention import Attention |
| 23 | from vllm.model_executor.layers.fused_moe import ( |
| 24 | FusedMoEFactory, |
| 25 | GateLinear, |
| 26 | fused_moe_make_expert_params_mapping, |
| 27 | ) |
| 28 | from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig |
| 29 | from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts |
| 30 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 31 | from vllm.model_executor.layers.linear import ( |
| 32 | QKVParallelLinear, |
| 33 | RowParallelLinear, |
| 34 | ) |
| 35 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 36 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 37 | from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_BLOCK_SIZE |
| 38 | from vllm.model_executor.layers.rotary_embedding import get_rope |
| 39 | from vllm.model_executor.layers.utils import rocm_unquantized_gemm |
| 40 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 41 | ParallelLMHead, |
| 42 | VocabParallelEmbedding, |
| 43 | ) |
| 44 | from vllm.model_executor.model_loader.weight_utils import ( |
| 45 | default_weight_loader, |
| 46 | maybe_remap_kv_scale_name, |
| 47 | remap_moe_expert_weights, |
| 48 | ) |
| 49 | from vllm.model_executor.models.utils import sequence_parallel_chunk |
| 50 | from vllm.platforms import current_platform |
| 51 | from vllm.sequence import IntermediateTensors |
| 52 | from vllm.utils.math_utils import cdiv |
| 53 | from vllm.v1.attention.backend import AttentionType |
| 54 | |
| 55 | from .interfaces import ( |
| 56 | EagleModelMixin, |
| 57 | SupportsEagle, |
| 58 | SupportsEagle3, |
| 59 | SupportsLoRA, |
| 60 | SupportsPP, |
| 61 | ) |
| 62 | from .utils import ( |
| 63 | AutoWeightsLoader, |
| 64 | WeightsMapper, |
| 65 | extract_layer_index, |
| 66 | is_pp_missing_parameter, |
| 67 | make_empty_intermediate_tensors_factory, |
| 68 | make_layers, |
| 69 | maybe_prefix, |
| 70 | ) |
| 71 | |
| 72 | _GPT_OSS_STREAMED_EXPERT_SUFFIX_TO_SHARD = { |
| 73 | "w13_weight": "gpt_oss_w13", |
| 74 | "w2_weight": "gpt_oss_w2", |
| 75 | "w13_bias": "gpt_oss_w13", |
| 76 | "w2_bias": "gpt_oss_w2", |
| 77 | "w13_weight_scale": "gpt_oss_w13", |
| 78 | "w2_weight_scale": "gpt_oss_w2", |
| 79 | "w13_weight_packed": "gpt_oss_w13", |
| 80 | "w2_weight_packed": "gpt_oss_w2", |
| 81 | } |
| 82 | |
| 83 | |
| 84 | def _get_weight_loader(param: torch.Tensor) -> Callable[..., object]: |
| 85 | weight_loader = getattr(param, "weight_loader", default_weight_loader) |
| 86 | if not callable(weight_loader): |
| 87 | raise TypeError("weight_loader must be callable") |
| 88 | return weight_loader |
| 89 | |
| 90 | |
| 91 | class OAIAttention(nn.Module): |
| 92 | # Override to switch RoPE convention. gpt-oss uses NeoX (chunk halves); |
| 93 | # privacy-filter and similar derivatives use GPT-J (interleaved pairs). |
| 94 | rope_is_neox_style: bool = True |
| 95 | |
| 96 | def __init__( |
| 97 | self, |
| 98 | config: GptOssConfig, |
| 99 | quant_config: QuantizationConfig | None = None, |
| 100 | cache_config: CacheConfig | None = None, |
| 101 | prefix: str = "", |
| 102 | ): |
| 103 | super().__init__() |
| 104 | self.layer_idx = extract_layer_index(prefix) |
| 105 | self.head_dim = config.head_dim |
| 106 | self.num_attention_heads = config.num_attention_heads |
| 107 | self.num_key_value_heads = config.num_key_value_heads |
| 108 | self.hidden_size = config.hidden_size |
| 109 | |
| 110 | self.rotary_emb = get_rope( |
| 111 | self.head_dim, |
| 112 | max_position=config.max_position_embeddings, |
| 113 | dtype=torch.float32, |
| 114 | rope_parameters={ |
| 115 | "rope_theta": config.rope_parameters["rope_theta"], |
| 116 | "rope_type": "yarn", |
| 117 | "factor": config.rope_parameters["factor"], |
| 118 | "original_max_position_embeddings": config.rope_parameters[ |
| 119 | "original_max_position_embeddings" |
| 120 | ], |
| 121 | "beta_fast": config.rope_parameters["beta_fast"], |
| 122 | "beta_slow": config.rope_parameters["beta_slow"], |
| 123 | "truncate": config.rope_parameters.get("truncate", True), |
| 124 | }, |
| 125 | is_neox_style=self.rope_is_neox_style, |
| 126 | ) |
| 127 | |
| 128 | tp_size = get_tensor_model_parallel_world_size() |
| 129 | |
| 130 | self.sinks = torch.nn.Parameter( |
| 131 | torch.empty(config.num_attention_heads // tp_size, requires_grad=False) |
| 132 | ) |
| 133 | |
| 134 | self.q_size = self.num_attention_heads * self.head_dim // tp_size |
| 135 | self.kv_size = self.num_key_value_heads * self.head_dim // tp_size |
| 136 | self.scaling = self.head_dim**-0.5 |
| 137 | |
| 138 | self.qkv_proj = QKVParallelLinear( |
| 139 | hidden_size=self.hidden_size, |
| 140 | head_size=self.head_dim, |
| 141 | total_num_heads=self.num_attention_heads, |
| 142 | total_num_kv_heads=self.num_key_value_heads, |
| 143 | bias=True, |
| 144 | quant_config=quant_config, |
| 145 | prefix=f"{prefix}.qkv_proj", |
| 146 | ) |
| 147 | |
| 148 | self.o_proj = RowParallelLinear( |
| 149 | input_size=self.num_attention_heads * self.head_dim, |
| 150 | output_size=self.hidden_size, |
| 151 | bias=True, |
| 152 | quant_config=quant_config, |
| 153 | prefix=f"{prefix}.o_proj", |
| 154 | ) |
| 155 | |
| 156 | self.num_local_attention_heads = config.num_attention_heads // tp_size |
| 157 | self.num_local_key_value_heads = config.num_key_value_heads // tp_size |
| 158 | |
| 159 | self.attn = self._build_attention( |
| 160 | config=config, |
| 161 | cache_config=cache_config, |
| 162 | quant_config=quant_config, |
| 163 | prefix=prefix, |
| 164 | ) |
| 165 | |
| 166 | def _build_attention( |
| 167 | self, |
| 168 | config: GptOssConfig, |
| 169 | cache_config: CacheConfig | None, |
| 170 | quant_config: QuantizationConfig | None, |
| 171 | prefix: str, |
| 172 | ) -> Attention: |
| 173 | # Override to swap in an encoder-only attention or alter the |
| 174 | # per-layer sliding-window policy. |
| 175 | # Only apply sliding window to every other layer |
| 176 | sliding_window = config.sliding_window if self.layer_idx % 2 == 0 else None |
| 177 | return Attention( |
| 178 | self.num_local_attention_heads, |
| 179 | self.head_dim, |
| 180 | self.scaling, |
| 181 | num_kv_heads=self.num_local_key_value_heads, |
| 182 | cache_config=cache_config, |
| 183 | quant_config=quant_config, |
| 184 | per_layer_sliding_window=sliding_window, |
| 185 | attn_type=AttentionType.DECODER, |
| 186 | prefix=f"{prefix}.attn", |
| 187 | sinks=self.sinks, |
| 188 | ) |
| 189 | |
| 190 | def forward( |
| 191 | self, hidden_states: torch.Tensor, positions: torch.Tensor |
| 192 | ) -> torch.Tensor: |
| 193 | qkv, _ = self.qkv_proj(hidden_states) |
| 194 | q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) |
| 195 | q, k = self.rotary_emb(positions, q, k) |
| 196 | attn_output = self.attn(q, k, v) |
| 197 | output, _ = self.o_proj(attn_output) |
| 198 | return output |
| 199 | |
| 200 | |
| 201 | class GptOssRoutedExperts(RoutedExperts): |
| 202 | """Load one GPT-OSS expert at a time without assembling the global stack.""" |
| 203 | |
| 204 | @staticmethod |
| 205 | def _narrow_for_rank( |
| 206 | loaded_weight: torch.Tensor, |
| 207 | dim: int, |
| 208 | rank: int, |
| 209 | size: int, |
| 210 | ) -> torch.Tensor: |
| 211 | start = rank * size |
| 212 | available = loaded_weight.shape[dim] - start |
| 213 | return loaded_weight.narrow(dim, start, min(size, max(available, 0))) |
| 214 | |
| 215 | @staticmethod |
| 216 | def _copy_to_expert(expert_data: torch.Tensor, loaded_weight: torch.Tensor) -> None: |
| 217 | if expert_data.numel() == loaded_weight.numel() == 1: |
| 218 | expert_data.copy_(loaded_weight.reshape(())) |
| 219 | return |
| 220 | while loaded_weight.ndim > expert_data.ndim and loaded_weight.shape[-1] == 1: |
| 221 | loaded_weight = loaded_weight.squeeze(-1) |
| 222 | slices = tuple(slice(0, size) for size in loaded_weight.shape) |
| 223 | expert_data[slices].copy_(loaded_weight) |
| 224 | |
| 225 | def _load_expert_bias( |
| 226 | self, |
| 227 | expert_data: torch.Tensor, |
| 228 | loaded_weight: torch.Tensor, |
| 229 | shard_id: str, |
| 230 | ) -> None: |
| 231 | tp_rank = self.moe_config.moe_parallel_config.tp_rank |
| 232 | if shard_id == "gpt_oss_w13": |
| 233 | loaded_weight = self._narrow_for_rank( |
| 234 | loaded_weight, 0, tp_rank, expert_data.shape[0] |
| 235 | ) |
| 236 | elif tp_rank != 0: |
| 237 | loaded_weight = torch.zeros_like(loaded_weight) |
| 238 | self._copy_to_expert(expert_data, loaded_weight) |
| 239 | |
| 240 | def _load_unquantized_expert( |
| 241 | self, |
| 242 | expert_data: torch.Tensor, |
| 243 | loaded_weight: torch.Tensor, |
| 244 | shard_id: str, |
| 245 | ) -> None: |
| 246 | tp_rank = self.moe_config.moe_parallel_config.tp_rank |
| 247 | if shard_id == "gpt_oss_w13": |
| 248 | loaded_weight = self._narrow_for_rank( |
| 249 | loaded_weight, 1, tp_rank, expert_data.shape[0] |
| 250 | ) |
| 251 | loaded_weight = loaded_weight.t().contiguous() |
| 252 | else: |
| 253 | loaded_weight = self._narrow_for_rank( |
| 254 | loaded_weight, 0, tp_rank, expert_data.shape[1] |
| 255 | ) |
| 256 | loaded_weight = loaded_weight.t().contiguous() |
| 257 | self._copy_to_expert(expert_data, loaded_weight) |
| 258 | |
| 259 | def _load_packed_expert( |
| 260 | self, |
| 261 | expert_data: torch.Tensor, |
| 262 | loaded_weight: torch.Tensor, |
| 263 | weight_name: str, |
| 264 | shard_id: str, |
| 265 | ) -> None: |
| 266 | tp_rank = self.moe_config.moe_parallel_config.tp_rank |
| 267 | is_w13 = shard_id == "gpt_oss_w13" |
| 268 | is_weight = weight_name.endswith(("_weight", "_weight_packed")) |
| 269 | is_partitioned_scale = weight_name.endswith("_weight_scale") |
| 270 | |
| 271 | if is_weight: |
| 272 | loaded_weight = loaded_weight.reshape(loaded_weight.shape[0], -1) |
| 273 | |
| 274 | if is_w13 and (is_weight or is_partitioned_scale): |
| 275 | loaded_weight = self._narrow_for_rank( |
| 276 | loaded_weight, 0, tp_rank, expert_data.shape[0] |
| 277 | ) |
| 278 | elif not is_w13 and (is_weight or is_partitioned_scale): |
| 279 | loaded_weight = self._narrow_for_rank( |
| 280 | loaded_weight, -1, tp_rank, expert_data.shape[-1] |
| 281 | ) |
| 282 | |
| 283 | self._copy_to_expert(expert_data, loaded_weight) |
| 284 | |
| 285 | @typing.overload |
| 286 | def weight_loader( |
| 287 | self, |
| 288 | param: torch.nn.Parameter, |
| 289 | loaded_weight: torch.Tensor, |
| 290 | weight_name: str, |
| 291 | shard_id: str, |
| 292 | expert_id: int, |
| 293 | return_success: typing.Literal[False], |
| 294 | ) -> None: ... |
| 295 | |
| 296 | @typing.overload |
| 297 | def weight_loader( |
| 298 | self, |
| 299 | param: torch.nn.Parameter, |
| 300 | loaded_weight: torch.Tensor, |
| 301 | weight_name: str, |
| 302 | shard_id: str, |
| 303 | expert_id: int, |
| 304 | return_success: typing.Literal[True], |
| 305 | ) -> bool: ... |
| 306 | |
| 307 | def weight_loader( |
| 308 | self, |
| 309 | param: torch.nn.Parameter, |
| 310 | loaded_weight: torch.Tensor, |
| 311 | weight_name: str, |
| 312 | shard_id: str, |
| 313 | expert_id: int, |
| 314 | return_success: bool = False, |
| 315 | ) -> bool | None: |
| 316 | if shard_id not in ("gpt_oss_w13", "gpt_oss_w2"): |
| 317 | if return_success: |
| 318 | return super().weight_loader( |
| 319 | param, |
| 320 | loaded_weight, |
| 321 | weight_name, |
| 322 | shard_id, |
| 323 | expert_id, |
| 324 | True, |
| 325 | ) |
| 326 | super().weight_loader( |
| 327 | param, |
| 328 | loaded_weight, |
| 329 | weight_name, |
| 330 | shard_id, |
| 331 | expert_id, |
| 332 | False, |
| 333 | ) |
| 334 | return None |
| 335 | |
| 336 | expert_id = self._map_global_expert_id_to_local_expert_id(expert_id) |
| 337 | if expert_id == -1: |
| 338 | return False if return_success else None |
| 339 | |
| 340 | expert_data = param.data[expert_id] |
| 341 | weight_dtype = getattr(self.quant_method, "weight_dtype", "") |
| 342 | quant_method_name = self.quant_method.__class__.__name__ |
| 343 | |
| 344 | if weight_name.endswith("_bias"): |
| 345 | self._load_expert_bias(expert_data, loaded_weight, shard_id) |
| 346 | elif weight_dtype in ( |
| 347 | "gpt_oss_mxfp4", |
| 348 | "mxfp4", |
| 349 | ) or quant_method_name in ( |
| 350 | "CompressedTensorsW4A4Nvfp4MoEMethod", |
| 351 | "Nvfp4OnlineMoEMethod", |
| 352 | "CompressedTensorsWNA16MoEMethod", |
| 353 | ): |
| 354 | self._load_packed_expert(expert_data, loaded_weight, weight_name, shard_id) |
| 355 | else: |
| 356 | self._load_unquantized_expert(expert_data, loaded_weight, shard_id) |
| 357 | return True if return_success else None |
| 358 | |
| 359 | |
| 360 | class MLPBlock(torch.nn.Module): |
| 361 | def __init__( |
| 362 | self, |
| 363 | vllm_config: VllmConfig, |
| 364 | layer_idx: int, |
| 365 | prefix: str = "", |
| 366 | ): |
| 367 | super().__init__() |
| 368 | |
| 369 | config = vllm_config.model_config.hf_config |
| 370 | quant_config = vllm_config.quant_config |
| 371 | parallel_config = vllm_config.parallel_config |
| 372 | |
| 373 | self.is_sequence_parallel = ( |
| 374 | parallel_config.use_sequence_parallel_moe |
| 375 | and vllm_config.lora_config is None |
| 376 | ) |
| 377 | |
| 378 | self.layer_idx = layer_idx |
| 379 | self.num_experts = config.num_local_experts |
| 380 | self.hidden_size = config.hidden_size |
| 381 | self.experts_per_token = config.num_experts_per_tok |
| 382 | self.world_size = dist.get_world_size() if dist.is_initialized() else 1 |
| 383 | self.router = GateLinear( |
| 384 | config.hidden_size, |
| 385 | config.num_local_experts, |
| 386 | bias=True, |
| 387 | prefix=f"{prefix}.router", |
| 388 | return_bias=False, |
| 389 | ) |
| 390 | assert config.intermediate_size % self.world_size == 0 |
| 391 | self.experts = FusedMoEFactory( |
| 392 | num_experts=config.num_local_experts, |
| 393 | top_k=config.num_experts_per_tok, |
| 394 | hidden_size=config.hidden_size, |
| 395 | intermediate_size=config.intermediate_size, |
| 396 | renormalize=True, |
| 397 | quant_config=quant_config, |
| 398 | prefix=f"{prefix}.experts", |
| 399 | apply_router_weight_on_input=False, |
| 400 | has_bias=True, |
| 401 | activation="swigluoai", |
| 402 | is_sequence_parallel=self.is_sequence_parallel, |
| 403 | routed_experts_cls=GptOssRoutedExperts, |
| 404 | ) |
| 405 | |
| 406 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 407 | num_tokens = x.shape[0] |
| 408 | if self.is_sequence_parallel: |
| 409 | x = sequence_parallel_chunk(x) |
| 410 | |
| 411 | if current_platform.is_rocm(): |
| 412 | g = rocm_unquantized_gemm( |
| 413 | self, x[:, : self.hidden_size], self.router.weight, self.router.bias |
| 414 | ) |
| 415 | else: |
| 416 | g = self.router(x) |
| 417 | x = self.experts(hidden_states=x, router_logits=g)[:, : self.hidden_size] |
| 418 | |
| 419 | if self.is_sequence_parallel: |
| 420 | x = tensor_model_parallel_all_gather(x.contiguous(), 0) |
| 421 | x = x[:num_tokens] |
| 422 | return x |
| 423 | |
| 424 | |
| 425 | class TransformerBlock(torch.nn.Module): |
| 426 | # Override to swap attention/MLP without re-implementing the block. |
| 427 | attention_cls: type[nn.Module] = OAIAttention |
| 428 | mlp_cls: type[nn.Module] = MLPBlock |
| 429 | |
| 430 | def __init__( |
| 431 | self, |
| 432 | vllm_config: VllmConfig, |
| 433 | quant_config: QuantizationConfig, |
| 434 | prefix: str = "", |
| 435 | ): |
| 436 | super().__init__() |
| 437 | |
| 438 | config = vllm_config.model_config.hf_config |
| 439 | cache_config = vllm_config.cache_config |
| 440 | |
| 441 | self.layer_idx = extract_layer_index(prefix) |
| 442 | self.attn = self.attention_cls( |
| 443 | config, |
| 444 | prefix=f"{prefix}.attn", |
| 445 | quant_config=quant_config, |
| 446 | cache_config=cache_config, |
| 447 | ) |
| 448 | self.mlp = self.mlp_cls(vllm_config, self.layer_idx, prefix=f"{prefix}.mlp") |
| 449 | self.input_layernorm = RMSNorm(config.hidden_size, eps=1e-5) |
| 450 | self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=1e-5) |
| 451 | |
| 452 | def forward( |
| 453 | self, |
| 454 | hidden_states: torch.Tensor, |
| 455 | positions: torch.Tensor, |
| 456 | residual: torch.Tensor | None, |
| 457 | ) -> torch.Tensor: |
| 458 | # Self Attention |
| 459 | if residual is None: |
| 460 | residual = hidden_states |
| 461 | hidden_states = self.input_layernorm(hidden_states) |
| 462 | else: |
| 463 | hidden_states, residual = self.input_layernorm(hidden_states, residual) |
| 464 | hidden_states = self.attn(hidden_states, positions) |
| 465 | |
| 466 | # Fully Connected |
| 467 | hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) |
| 468 | output = self.mlp(hidden_states) |
| 469 | return output, residual |
| 470 | |
| 471 | |
| 472 | @support_torch_compile |
| 473 | class GptOssModel(nn.Module, EagleModelMixin): |
| 474 | # Override to swap in an alternative TransformerBlock subclass. |
| 475 | block_cls: type[nn.Module] = TransformerBlock |
| 476 | |
| 477 | def __init__( |
| 478 | self, |
| 479 | *, |
| 480 | vllm_config: VllmConfig, |
| 481 | prefix: str = "", |
| 482 | ): |
| 483 | super().__init__() |
| 484 | self.config = vllm_config.model_config.hf_config |
| 485 | self.quant_config = vllm_config.quant_config |
| 486 | self.parallel_config = vllm_config.parallel_config |
| 487 | self.embedding = VocabParallelEmbedding( |
| 488 | self.config.vocab_size, |
| 489 | self.config.hidden_size, |
| 490 | ) |
| 491 | self.start_layer, self.end_layer, self.layers = make_layers( |
| 492 | self.config.num_hidden_layers, |
| 493 | lambda prefix: self.block_cls( |
| 494 | vllm_config, |
| 495 | prefix=prefix, |
| 496 | quant_config=self.quant_config, |
| 497 | ), |
| 498 | prefix=f"{prefix}.layers", |
| 499 | ) |
| 500 | self.norm = RMSNorm(self.config.hidden_size, eps=1e-5) |
| 501 | self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( |
| 502 | ["hidden_states", "residual"], self.config.hidden_size |
| 503 | ) |
| 504 | |
| 505 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 506 | return self.embedding(input_ids) |
| 507 | |
| 508 | def forward( |
| 509 | self, |
| 510 | input_ids: torch.Tensor | None, |
| 511 | positions: torch.Tensor, |
| 512 | intermediate_tensors: IntermediateTensors | None = None, |
| 513 | inputs_embeds: torch.Tensor | None = None, |
| 514 | ) -> torch.Tensor: |
| 515 | if get_pp_group().is_first_rank: |
| 516 | if inputs_embeds is not None: |
| 517 | x = inputs_embeds |
| 518 | else: |
| 519 | x = self.embed_input_ids(input_ids) |
| 520 | |
| 521 | residual = None |
| 522 | else: |
| 523 | assert intermediate_tensors is not None |
| 524 | x = intermediate_tensors["hidden_states"] |
| 525 | residual = intermediate_tensors["residual"] |
| 526 | |
| 527 | aux_hidden_states = self._maybe_add_hidden_state( |
| 528 | [], self.start_layer, x, residual |
| 529 | ) |
| 530 | for i in range(self.start_layer, self.end_layer): |
| 531 | layer = self.layers[i] |
| 532 | x, residual = layer(x, positions, residual) |
| 533 | self._maybe_add_hidden_state(aux_hidden_states, i + 1, x, residual) |
| 534 | if not get_pp_group().is_last_rank: |
| 535 | return IntermediateTensors({"hidden_states": x, "residual": residual}) |
| 536 | x, _ = self.norm(x, residual) |
| 537 | |
| 538 | if len(aux_hidden_states) > 0: |
| 539 | return x, aux_hidden_states |
| 540 | return x |
| 541 | |
| 542 | def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: |
| 543 | # Params for weights, weight scales, activation scales |
| 544 | # (param_name, weight_name, expert_id, shard_id) |
| 545 | # NOTE: this is only used for quark. |
| 546 | return fused_moe_make_expert_params_mapping( |
| 547 | self, |
| 548 | ckpt_gate_proj_name="w1", |
| 549 | ckpt_down_proj_name="w2", |
| 550 | ckpt_up_proj_name="w3", |
| 551 | num_experts=self.config.num_local_experts, |
| 552 | num_redundant_experts=0, |
| 553 | ) |
| 554 | |
| 555 | def _load_weights_mxfp4( |
| 556 | self, |
| 557 | ep_rank_end: int, |
| 558 | ep_rank_start: int, |
| 559 | heads_per_rank: int, |
| 560 | head_start: int, |
| 561 | weights: Iterable[tuple[str, torch.Tensor]], |
| 562 | stacked_params_mapping: list[tuple[str, str, str]], |
| 563 | ) -> set[str]: |
| 564 | params_dict = dict(self.named_parameters()) |
| 565 | loaded_params: set[str] = set() |
| 566 | |
| 567 | use_ep = self.parallel_config.enable_expert_parallel |
| 568 | num_experts = self.config.num_local_experts |
| 569 | |
| 570 | # In MoE, we need to flatten the tensor parallel size across the data |
| 571 | # parallel size when EP is disabled. |
| 572 | tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp( |
| 573 | tp_size=get_tensor_model_parallel_world_size(), |
| 574 | dp_size=get_dp_group().world_size, |
| 575 | dp_rank=get_dp_group().rank_in_group, |
| 576 | pcp_size=get_pcp_group().world_size, |
| 577 | pcp_rank=get_pcp_group().rank_in_group, |
| 578 | ) |
| 579 | |
| 580 | intermediate_size = self.config.intermediate_size |
| 581 | intermediate_size_block = intermediate_size // OCP_MX_BLOCK_SIZE |
| 582 | per_rank_intermediate_size_block = cdiv(intermediate_size_block, tp_size) |
| 583 | per_rank_intermediate_size = ( |
| 584 | per_rank_intermediate_size_block * OCP_MX_BLOCK_SIZE |
| 585 | ) |
| 586 | |
| 587 | # Calculate common slicing bounds for current rank |
| 588 | tp_rank_start = tp_rank * per_rank_intermediate_size |
| 589 | tp_rank_end = min((tp_rank + 1) * per_rank_intermediate_size, intermediate_size) |
| 590 | |
| 591 | # Use centralized weight remapping for MoE expert parameters |
| 592 | for name, weight in remap_moe_expert_weights(weights, params_dict): |
| 593 | # Skip layers on other devices. |
| 594 | if is_pp_missing_parameter(name, self): |
| 595 | continue |
| 596 | |
| 597 | if self._try_load_streamed_expert(name, weight, params_dict, loaded_params): |
| 598 | continue |
| 599 | |
| 600 | if ".w13_weight_scale" in name: |
| 601 | # Handle MLP gate and up projection weights scale |
| 602 | if use_ep: |
| 603 | narrow_weight = weight[ep_rank_start:ep_rank_end, ...] |
| 604 | else: |
| 605 | narrow_weight = weight[:, 2 * tp_rank_start : 2 * tp_rank_end, ...] |
| 606 | |
| 607 | param = params_dict[name] |
| 608 | weight_loader = _get_weight_loader(param) |
| 609 | weight_loader( |
| 610 | param, |
| 611 | narrow_weight, |
| 612 | weight_name=name, |
| 613 | shard_id=None, |
| 614 | expert_id=None, |
| 615 | ) |
| 616 | loaded_params.add(name) |
| 617 | continue |
| 618 | elif ".w2_weight_scale" in name: |
| 619 | # Handle MLP down projection weights |
| 620 | if use_ep: |
| 621 | narrow_weight = weight[ep_rank_start:ep_rank_end, ...] |
| 622 | else: |
| 623 | narrow_weight = weight[ |
| 624 | ..., |
| 625 | tp_rank_start // OCP_MX_BLOCK_SIZE : tp_rank_end |
| 626 | // OCP_MX_BLOCK_SIZE, |
| 627 | ] |
| 628 | |
| 629 | param = params_dict[name] |
| 630 | weight_loader = _get_weight_loader(param) |
| 631 | weight_loader( |
| 632 | param, |
| 633 | narrow_weight, |
| 634 | weight_name=name, |
| 635 | shard_id=None, |
| 636 | expert_id=None, |
| 637 | ) |
| 638 | loaded_params.add(name) |
| 639 | continue |
| 640 | elif ".w13_weight" in name: |
| 641 | # Handle MLP gate and up projection weights |
| 642 | # flat weight from (E, 2 * N, block_size, entry_per_block) |
| 643 | # to (E, 2 * N, -1), shouldn't trigger copy for contiguous |
| 644 | weight = weight.view( |
| 645 | num_experts, 2 * intermediate_size, -1 |
| 646 | ).contiguous() |
| 647 | |
| 648 | # Extract gate and up projection parts |
| 649 | # since the weight is shuffled, we can slice directly |
| 650 | if use_ep: |
| 651 | narrow_weight = weight[ep_rank_start:ep_rank_end, ...] |
| 652 | else: |
| 653 | narrow_weight = weight[:, 2 * tp_rank_start : 2 * tp_rank_end, ...] |
| 654 | |
| 655 | param = params_dict[name] |
| 656 | weight_loader = _get_weight_loader(param) |
| 657 | weight_loader( |
| 658 | param, |
| 659 | narrow_weight, |
| 660 | weight_name=name, |
| 661 | shard_id=None, |
| 662 | expert_id=None, |
| 663 | ) |
| 664 | loaded_params.add(name) |
| 665 | continue |
| 666 | elif ".w2_weight" in name: |
| 667 | # Handle MLP down projection weights |
| 668 | # same flatten here, but since 2 mx4 value are packed in 1 |
| 669 | # uint8, divide by 2 |
| 670 | weight = weight.view( |
| 671 | num_experts, -1, intermediate_size // 2 |
| 672 | ).contiguous() |
| 673 | if use_ep: |
| 674 | narrow_weight = weight[ep_rank_start:ep_rank_end, ...] |
| 675 | else: |
| 676 | narrow_weight = weight[..., tp_rank_start // 2 : tp_rank_end // 2] |
| 677 | |
| 678 | param = params_dict[name] |
| 679 | weight_loader = _get_weight_loader(param) |
| 680 | weight_loader( |
| 681 | param, |
| 682 | narrow_weight, |
| 683 | weight_name=name, |
| 684 | shard_id=None, |
| 685 | expert_id=None, |
| 686 | ) |
| 687 | loaded_params.add(name) |
| 688 | continue |
| 689 | elif ".w13_bias" in name: |
| 690 | # Handle MLP gate and up projection biases |
| 691 | # Extract gate and up projection bias parts |
| 692 | if use_ep: |
| 693 | narrow_weight = weight[ep_rank_start:ep_rank_end, ...] |
| 694 | else: |
| 695 | narrow_weight = weight[:, 2 * tp_rank_start : 2 * tp_rank_end] |
| 696 | |
| 697 | param = params_dict[name] |
| 698 | weight_loader = _get_weight_loader(param) |
| 699 | weight_loader( |
| 700 | param, |
| 701 | narrow_weight, |
| 702 | weight_name=name, |
| 703 | shard_id=None, |
| 704 | expert_id=None, |
| 705 | ) |
| 706 | loaded_params.add(name) |
| 707 | continue |
| 708 | elif ".w2_bias" in name: |
| 709 | # Handle MLP down projection bias |
| 710 | param = params_dict[name] |
| 711 | weight_loader = _get_weight_loader(param) |
| 712 | if use_ep: |
| 713 | weight = weight[ep_rank_start:ep_rank_end, ...] |
| 714 | else: |
| 715 | # (only load on rank 0 to avoid duplication) |
| 716 | if tp_rank != 0: |
| 717 | weight.zero_() |
| 718 | weight_loader( |
| 719 | param, weight, weight_name=name, shard_id=None, expert_id=None |
| 720 | ) |
| 721 | loaded_params.add(name) |
| 722 | continue |
| 723 | elif "sinks" in name: |
| 724 | # Handle attention sinks (distributed across ranks) |
| 725 | param = params_dict[name] |
| 726 | narrow_weight = weight.narrow(0, head_start, heads_per_rank) |
| 727 | param.data.copy_(narrow_weight) |
| 728 | loaded_params.add(name) |
| 729 | continue |
| 730 | for param_name, weight_name, shard_id in stacked_params_mapping: |
| 731 | if weight_name not in name: |
| 732 | continue |
| 733 | name = name.replace(weight_name, param_name) |
| 734 | param = params_dict[name] |
| 735 | weight_loader = _get_weight_loader(param) |
| 736 | if weight_loader == default_weight_loader: |
| 737 | weight_loader(param, weight) |
| 738 | else: |
| 739 | weight_loader(param, weight, shard_id) |
| 740 | break |
| 741 | else: |
| 742 | # Handle all other weights with potential renaming |
| 743 | if name not in params_dict: |
| 744 | continue |
| 745 | param = params_dict[name] |
| 746 | weight_loader = _get_weight_loader(param) |
| 747 | weight_loader(param, weight) |
| 748 | loaded_params.add(name) |
| 749 | return loaded_params |
| 750 | |
| 751 | def _load_weights_quark( |
| 752 | self, |
| 753 | ep_rank_end: int, |
| 754 | ep_rank_start: int, |
| 755 | heads_per_rank: int, |
| 756 | head_start: int, |
| 757 | weights: Iterable[tuple[str, torch.Tensor]], |
| 758 | stacked_params_mapping: list[tuple[str, str, str]], |
| 759 | ) -> set[str]: |
| 760 | params_dict = dict(self.named_parameters()) |
| 761 | loaded_params: set[str] = set() |
| 762 | |
| 763 | use_ep = self.parallel_config.enable_expert_parallel |
| 764 | num_experts = self.config.num_local_experts |
| 765 | |
| 766 | if use_ep: |
| 767 | tp_rank = get_tensor_model_parallel_rank() |
| 768 | tp_size = get_tensor_model_parallel_world_size() |
| 769 | else: |
| 770 | tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp( |
| 771 | tp_size=get_tensor_model_parallel_world_size(), |
| 772 | dp_size=get_dp_group().world_size, |
| 773 | dp_rank=get_dp_group().rank_in_group, |
| 774 | pcp_size=get_pcp_group().world_size, |
| 775 | pcp_rank=get_pcp_group().rank_in_group, |
| 776 | ) |
| 777 | |
| 778 | def _is_mxfp4(weight_dtype: str | None) -> bool: |
| 779 | """Return True for any MXFP4 weight-dtype variant. |
| 780 | |
| 781 | Covers "gpt_oss_mxfp4" (GptOssMxfp4MoEMethod) and "mxfp4" |
| 782 | (QuarkMoEMethod with fp4 weights) and any future variants. |
| 783 | """ |
| 784 | return weight_dtype is not None and "mxfp4" in weight_dtype |
| 785 | |
| 786 | def _get_moe_weight_dtype(layer_id: int = 0) -> str | None: |
| 787 | """Helper function to get MoE quantization weight dtype. |
| 788 | |
| 789 | Args: |
| 790 | layer_id: Layer index to check (default 0, as all layers should |
| 791 | have the same quantization method) |
| 792 | |
| 793 | Returns: |
| 794 | Weight dtype string (e.g., "mxfp4", "fp8") or None if not available |
| 795 | |
| 796 | """ |
| 797 | if hasattr(self.layers[layer_id].mlp.experts._quant_method, "weight_dtype"): |
| 798 | return self.layers[layer_id].mlp.experts._quant_method.weight_dtype |
| 799 | return None |
| 800 | |
| 801 | intermediate_size = self.config.intermediate_size |
| 802 | |
| 803 | moe_weight_dtype = _get_moe_weight_dtype(layer_id=0) |
| 804 | |
| 805 | if _is_mxfp4(moe_weight_dtype): |
| 806 | # MXFP4 requires OCP_MX_BLOCK_SIZE alignment |
| 807 | intermediate_size_block = intermediate_size // OCP_MX_BLOCK_SIZE |
| 808 | per_rank_intermediate_size_block = cdiv(intermediate_size_block, tp_size) |
| 809 | per_rank_intermediate_size = ( |
| 810 | per_rank_intermediate_size_block * OCP_MX_BLOCK_SIZE |
| 811 | ) |
| 812 | else: |
| 813 | # FP8 and other formats don't need alignment |
| 814 | per_rank_intermediate_size = cdiv(intermediate_size, tp_size) |
| 815 | |
| 816 | tp_rank_start = tp_rank * per_rank_intermediate_size |
| 817 | tp_rank_end = min((tp_rank + 1) * per_rank_intermediate_size, intermediate_size) |
| 818 | expert_params_mapping = self.get_expert_mapping() |
| 819 | # Streamed per-expert reload is intentionally unsupported for Quark. |
| 820 | for name, loaded_weight in weights: |
| 821 | if is_pp_missing_parameter(name, self): |
| 822 | continue |
| 823 | |
| 824 | layer_id, expert_id, fused_name = None, None, None |
| 825 | moe_quant_method = None |
| 826 | if "experts" in name: |
| 827 | parts = name.split(".") |
| 828 | ids = [s for s in parts if s.isdigit()] |
| 829 | |
| 830 | # for amd-quark format that each expert is separated |
| 831 | # need to extract the parameter name with experts fused. |
| 832 | # example model: amd/gpt-oss-20b-MoE-Quant-W-MXFP4-A-FP8-KV-FP8 |
| 833 | if len(ids) == 2: |
| 834 | layer_id, expert_id = int(ids[0]), int(ids[-1]) |
| 835 | parts.pop(len(parts) - 1 - parts[::-1].index(str(expert_id))) |
| 836 | fused_name = ".".join(parts) |
| 837 | |
| 838 | # for openai mxfp4 format that all experts are combined |
| 839 | # no need to extract the parameter name with experts fused. |
| 840 | # models: openai/gpt-oss-20b, openai/gpt-oss-120b |
| 841 | elif len(ids) == 1: |
| 842 | layer_id, expert_id = int(ids[0]), None |
| 843 | fused_name = name |
| 844 | |
| 845 | else: |
| 846 | raise NameError( |
| 847 | f"Layer {name} contains more than 2 numeric indices. This is " |
| 848 | "an unexpected condition. Please open an issue if encountered." |
| 849 | ) |
| 850 | |
| 851 | # The MoE refactor (#41184) moved expert params under |
| 852 | # `mlp.experts.routed_experts.*`; remap the legacy checkpoint |
| 853 | # name so keys like w2_bias resolve against params_dict. |
| 854 | fused_name = fused_name.replace( |
| 855 | ".mlp.experts.", ".mlp.experts.routed_experts." |
| 856 | ) |
| 857 | |
| 858 | moe_quant_method = _get_moe_weight_dtype(layer_id=layer_id) |
| 859 | |
| 860 | if ( |
| 861 | all(key in name for key in ["input_scale", "mlp.experts"]) |
| 862 | and expert_id is not None |
| 863 | ): |
| 864 | assert loaded_weight.numel() == 1 |
| 865 | assert fused_name is not None |
| 866 | expert_data = params_dict[fused_name].data[expert_id] |
| 867 | expert_data.copy_(loaded_weight) |
| 868 | loaded_params.add(fused_name) |
| 869 | continue |
| 870 | |
| 871 | # Unified handler for mxfp4 weights and scales |
| 872 | elif _is_mxfp4(moe_quant_method) and any( |
| 873 | name.endswith(suffix) |
| 874 | for suffix in [ |
| 875 | ".w13_weight_scale", |
| 876 | ".w2_weight_scale", |
| 877 | ".w13_weight", |
| 878 | ".w2_weight", |
| 879 | ] |
| 880 | ): |
| 881 | is_w13 = ".w13_" in name |
| 882 | is_scale = "_scale" in name |
| 883 | |
| 884 | # Reshape weight for mxfp4 if needed (not for scales) |
| 885 | if not is_scale and expert_id is None: |
| 886 | if is_w13: |
| 887 | if loaded_weight.dim() < 3: |
| 888 | raise ValueError( |
| 889 | f"Expected w13_weight to have at least 3 " |
| 890 | f"dimensions, got shape " |
| 891 | f"{loaded_weight.shape}" |
| 892 | ) |
| 893 | if loaded_weight.shape[0] != num_experts: |
| 894 | raise ValueError( |
| 895 | f"Expected w13_weight first dimension to be " |
| 896 | f"{num_experts}, got " |
| 897 | f"{loaded_weight.shape[0]}" |
| 898 | ) |
| 899 | loaded_weight = loaded_weight.view( |
| 900 | num_experts, 2 * intermediate_size, -1 |
| 901 | ).contiguous() |
| 902 | else: |
| 903 | if loaded_weight.dim() < 3: |
| 904 | raise ValueError( |
| 905 | f"Expected w2_weight to have at least 3 " |
| 906 | f"dimensions, got shape " |
| 907 | f"{loaded_weight.shape}" |
| 908 | ) |
| 909 | if loaded_weight.shape[0] != num_experts: |
| 910 | raise ValueError( |
| 911 | f"Expected w2_weight first dimension to be " |
| 912 | f"{num_experts}, got " |
| 913 | f"{loaded_weight.shape[0]}" |
| 914 | ) |
| 915 | loaded_weight = loaded_weight.view( |
| 916 | num_experts, -1, intermediate_size // 2 |
| 917 | ).contiguous() |
| 918 | |
| 919 | if use_ep: |
| 920 | sliced_weight = loaded_weight[ep_rank_start:ep_rank_end, ...] |
| 921 | else: |
| 922 | if is_w13: |
| 923 | if expert_id is None: |
| 924 | sliced_weight = loaded_weight[ |
| 925 | :, 2 * tp_rank_start : 2 * tp_rank_end, ... |
| 926 | ] |
| 927 | else: |
| 928 | sliced_weight = loaded_weight[ |
| 929 | 2 * tp_rank_start : 2 * tp_rank_end, ... |
| 930 | ] |
| 931 | else: |
| 932 | if is_scale: |
| 933 | sliced_weight = loaded_weight[ |
| 934 | ..., |
| 935 | tp_rank_start // OCP_MX_BLOCK_SIZE : tp_rank_end |
| 936 | // OCP_MX_BLOCK_SIZE, |
| 937 | ] |
| 938 | else: |
| 939 | sliced_weight = loaded_weight[ |
| 940 | ..., tp_rank_start // 2 : tp_rank_end // 2 |
| 941 | ] |
| 942 | |
| 943 | # NOTE(rob): because gpt-oss ckpt has "unique" structure with |
| 944 | # fused gate_up_proj fused on disk, we cannot use the existing |
| 945 | # weight loaders without added complexity, so just do the |
| 946 | # direct load here. |
| 947 | assert fused_name is not None |
| 948 | param = params_dict[fused_name] |
| 949 | expert_data = param.data[expert_id] |
| 950 | dim1 = sliced_weight.shape[0] |
| 951 | dim2 = sliced_weight.shape[1] |
| 952 | expert_data.data[:dim1, :dim2].copy_(sliced_weight) |
| 953 | loaded_params.add(fused_name) |
| 954 | continue |
| 955 | |
| 956 | elif name.endswith(".w13_weight") and moe_quant_method == "fp8": |
| 957 | if use_ep: |
| 958 | narrow_weight = loaded_weight[ep_rank_start:ep_rank_end, ...] |
| 959 | else: |
| 960 | if expert_id is None: |
| 961 | narrow_weight = loaded_weight[ |
| 962 | :, 2 * tp_rank_start : 2 * tp_rank_end, : |
| 963 | ] |
| 964 | else: |
| 965 | narrow_weight = loaded_weight[ |
| 966 | 2 * tp_rank_start : 2 * tp_rank_end, : |
| 967 | ] |
| 968 | |
| 969 | assert fused_name is not None |
| 970 | param = params_dict[fused_name] |
| 971 | |
| 972 | if expert_id is None: |
| 973 | param.data.copy_(narrow_weight) |
| 974 | else: |
| 975 | param.data[expert_id].copy_(narrow_weight) |
| 976 | |
| 977 | loaded_params.add(fused_name) |
| 978 | continue |
| 979 | |
| 980 | elif name.endswith(".w13_weight_scale") and moe_quant_method == "fp8": |
| 981 | assert fused_name is not None |
| 982 | param = params_dict[fused_name] |
| 983 | |
| 984 | # Check if this is per-channel or per-tensor scale |
| 985 | if loaded_weight.numel() > 1 and loaded_weight.dim() == 1: |
| 986 | if use_ep: |
| 987 | narrow_weight = loaded_weight[ep_rank_start:ep_rank_end, ...] |
| 988 | else: |
| 989 | narrow_weight = loaded_weight[ |
| 990 | 2 * tp_rank_start : 2 * tp_rank_end |
| 991 | ] |
| 992 | else: |
| 993 | narrow_weight = loaded_weight |
| 994 | |
| 995 | if expert_id is None: |
| 996 | param.data.copy_(narrow_weight) |
| 997 | else: |
| 998 | param.data[expert_id].copy_(narrow_weight) |
| 999 | |
| 1000 | loaded_params.add(fused_name) |
| 1001 | continue |
| 1002 | |
| 1003 | elif name.endswith(".w13_input_scale") and moe_quant_method == "fp8": |
| 1004 | assert fused_name is not None |
| 1005 | param = params_dict[fused_name] |
| 1006 | |
| 1007 | if expert_id is None: |
| 1008 | param.data.copy_(loaded_weight) |
| 1009 | else: |
| 1010 | param.data[expert_id].copy_(loaded_weight) |
| 1011 | |
| 1012 | loaded_params.add(fused_name) |
| 1013 | continue |
| 1014 | |
| 1015 | elif name.endswith(".w2_weight") and moe_quant_method == "fp8": |
| 1016 | if use_ep: |
| 1017 | narrow_weight = loaded_weight[ep_rank_start:ep_rank_end, ...] |
| 1018 | else: |
| 1019 | if expert_id is None: |
| 1020 | narrow_weight = loaded_weight[..., tp_rank_start:tp_rank_end] |
| 1021 | else: |
| 1022 | narrow_weight = loaded_weight[..., tp_rank_start:tp_rank_end] |
| 1023 | |
| 1024 | assert fused_name is not None |
| 1025 | param = params_dict[fused_name] |
| 1026 | |
| 1027 | if expert_id is None: |
| 1028 | param.data.copy_(narrow_weight) |
| 1029 | else: |
| 1030 | param.data[expert_id].copy_(narrow_weight) |
| 1031 | |
| 1032 | loaded_params.add(fused_name) |
| 1033 | continue |
| 1034 | |
| 1035 | elif name.endswith(".w2_weight_scale") and moe_quant_method == "fp8": |
| 1036 | assert fused_name is not None |
| 1037 | param = params_dict[fused_name] |
| 1038 | |
| 1039 | if use_ep: |
| 1040 | narrow_weight = loaded_weight[ep_rank_start:ep_rank_end, ...] |
| 1041 | else: |
| 1042 | narrow_weight = loaded_weight |
| 1043 | |
| 1044 | if expert_id is None: |
| 1045 | param.data.copy_(narrow_weight) |
| 1046 | else: |
| 1047 | param.data[expert_id].copy_(narrow_weight) |
| 1048 | |
| 1049 | loaded_params.add(fused_name) |
| 1050 | continue |
| 1051 | |
| 1052 | # Unified handler for bias loading (w13_bias and w2_bias) |
| 1053 | elif name.endswith(".w13_bias") or name.endswith(".w2_bias"): |
| 1054 | is_w13_bias = name.endswith(".w13_bias") |
| 1055 | |
| 1056 | if use_ep: |
| 1057 | sliced_weight = loaded_weight[ep_rank_start:ep_rank_end, ...] |
| 1058 | else: |
| 1059 | if is_w13_bias: |
| 1060 | if expert_id is None: |
| 1061 | sliced_weight = loaded_weight[ |
| 1062 | :, 2 * tp_rank_start : 2 * tp_rank_end |
| 1063 | ] |
| 1064 | else: |
| 1065 | sliced_weight = loaded_weight[ |
| 1066 | 2 * tp_rank_start : 2 * tp_rank_end |
| 1067 | ] |
| 1068 | else: |
| 1069 | sliced_weight = loaded_weight |
| 1070 | if tp_rank != 0: |
| 1071 | sliced_weight = sliced_weight.zero_() |
| 1072 | |
| 1073 | # NOTE(rob): because gpt-oss ckpt has "unique" structure with |
| 1074 | # fused gate_up_proj fused on disk, we cannot use the existing |
| 1075 | # weight loaders without added complexity, so just do the |
| 1076 | # direct load here. |
| 1077 | assert fused_name is not None |
| 1078 | param = params_dict[fused_name] |
| 1079 | expert_data = param.data[expert_id] |
| 1080 | dim1 = sliced_weight.shape[0] |
| 1081 | expert_data.data[:dim1].copy_(sliced_weight) |
| 1082 | loaded_params.add(fused_name) |
| 1083 | continue |
| 1084 | |
| 1085 | elif "sinks" in name: |
| 1086 | # Handle attention sinks (distributed across ranks) |
| 1087 | param = params_dict[name] |
| 1088 | narrow_weight = loaded_weight.narrow(0, head_start, heads_per_rank) |
| 1089 | param.data.copy_(narrow_weight) |
| 1090 | loaded_params.add(name) |
| 1091 | continue |
| 1092 | |
| 1093 | for param_name, weight_name, shard_id in stacked_params_mapping: |
| 1094 | # Skip non-stacked layers and experts (experts handled below). |
| 1095 | if weight_name not in name: |
| 1096 | continue |
| 1097 | # We have mlp.experts[0].gate_proj in the checkpoint. |
| 1098 | # Since we handle the experts below in expert_params_mapping, |
| 1099 | # we need to skip here BEFORE we update the name, otherwise |
| 1100 | # name will be updated to mlp.experts[0].gate_up_proj, which |
| 1101 | # will then be updated below in expert_params_mapping |
| 1102 | # for mlp.experts[0].gate_gate_up_proj, which breaks load. |
| 1103 | if ("mlp.experts." in name) and name not in params_dict: |
| 1104 | continue |
| 1105 | name = name.replace(weight_name, param_name) |
| 1106 | |
| 1107 | if name.endswith("scale"): |
| 1108 | # Remapping the name of FP8 kv-scale. |
| 1109 | remapped_name = maybe_remap_kv_scale_name(name, params_dict) |
| 1110 | if remapped_name is None: |
| 1111 | continue |
| 1112 | name = remapped_name |
| 1113 | |
| 1114 | param = params_dict[name] |
| 1115 | weight_loader = param.weight_loader |
| 1116 | |
| 1117 | weight_loader(param, loaded_weight, shard_id) |
| 1118 | loaded_params.add(name) |
| 1119 | break |
| 1120 | else: |
| 1121 | for mapping in expert_params_mapping: |
| 1122 | # Anyway, this is an expert weight and should not be |
| 1123 | # attempted to load as other weights later |
| 1124 | param_name, weight_name, mapping_expert_id, shard_id = mapping |
| 1125 | weight_name = ( |
| 1126 | weight_name[:-1] if weight_name.endswith(".") else weight_name |
| 1127 | ) |
| 1128 | |
| 1129 | if weight_name not in name: |
| 1130 | continue |
| 1131 | |
| 1132 | assert fused_name is not None |
| 1133 | param = params_dict[fused_name] |
| 1134 | # We should ask the weight loader to return success or not |
| 1135 | # here since otherwise we may skip experts with other |
| 1136 | # available replicas. |
| 1137 | weight_loader = typing.cast( |
| 1138 | Callable[..., bool], param.weight_loader |
| 1139 | ) |
| 1140 | # Use checkpoint's expert_id for quark format (when expert_id |
| 1141 | # is extracted from weight name), otherwise use mapping's expert_id |
| 1142 | actual_expert_id = ( |
| 1143 | expert_id if expert_id is not None else mapping_expert_id |
| 1144 | ) |
| 1145 | success = weight_loader( |
| 1146 | param, |
| 1147 | loaded_weight, |
| 1148 | fused_name, |
| 1149 | shard_id=shard_id, |
| 1150 | expert_id=actual_expert_id, |
| 1151 | return_success=True, |
| 1152 | ) |
| 1153 | if success: |
| 1154 | name = fused_name |
| 1155 | loaded_params.add(name) |
| 1156 | break |
| 1157 | else: |
| 1158 | if name not in params_dict: |
| 1159 | continue |
| 1160 | param = params_dict[name] |
| 1161 | weight_loader = _get_weight_loader(param) |
| 1162 | weight_loader(param, loaded_weight) |
| 1163 | |
| 1164 | loaded_params.add(name) |
| 1165 | return loaded_params |
| 1166 | |
| 1167 | def _load_weights_other( |
| 1168 | self, |
| 1169 | ep_rank_end: int, |
| 1170 | ep_rank_start: int, |
| 1171 | heads_per_rank: int, |
| 1172 | head_start: int, |
| 1173 | weights: Iterable[tuple[str, torch.Tensor]], |
| 1174 | stacked_params_mapping: list[tuple[str, str, str]], |
| 1175 | ) -> set[str]: |
| 1176 | params_dict = dict(self.named_parameters()) |
| 1177 | loaded_params: set[str] = set() |
| 1178 | |
| 1179 | use_ep = self.parallel_config.enable_expert_parallel |
| 1180 | |
| 1181 | # In MoE, we need to flatten the tensor parallel size across the data |
| 1182 | # parallel size when EP is disabled. |
| 1183 | tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp( |
| 1184 | tp_size=get_tensor_model_parallel_world_size(), |
| 1185 | dp_size=get_dp_group().world_size, |
| 1186 | dp_rank=get_dp_group().rank_in_group, |
| 1187 | pcp_size=get_pcp_group().world_size, |
| 1188 | pcp_rank=get_pcp_group().rank_in_group, |
| 1189 | ) |
| 1190 | |
| 1191 | intermediate_size = self.config.intermediate_size |
| 1192 | per_rank_intermediate_size = cdiv(intermediate_size, tp_size) |
| 1193 | # Calculate common slicing bounds for current rank |
| 1194 | tp_rank_start = tp_rank * per_rank_intermediate_size |
| 1195 | tp_rank_end = min((tp_rank + 1) * per_rank_intermediate_size, intermediate_size) |
| 1196 | |
| 1197 | # Use centralized weight remapping for MoE expert parameters. |
| 1198 | # The MoERunner refactor moved expert params under |
| 1199 | # `mlp.experts.routed_experts.*`; this remaps checkpoint names so |
| 1200 | # MoE weight/bias keys resolve against params_dict. |
| 1201 | for name, weight in remap_moe_expert_weights(weights, params_dict): |
| 1202 | # Skip layers on other devices. |
| 1203 | if is_pp_missing_parameter(name, self): |
| 1204 | continue |
| 1205 | |
| 1206 | if self._try_load_streamed_expert(name, weight, params_dict, loaded_params): |
| 1207 | continue |
| 1208 | |
| 1209 | if ".w13_weight" in name: |
| 1210 | # Handle MLP gate and up projection weights |
| 1211 | # Extract gate and up projection parts |
| 1212 | if use_ep: |
| 1213 | narrow_weight = weight[ep_rank_start:ep_rank_end, ...] |
| 1214 | else: |
| 1215 | narrow_weight = weight[:, :, 2 * tp_rank_start : 2 * tp_rank_end] |
| 1216 | |
| 1217 | narrow_weight = narrow_weight.permute(0, 2, 1).contiguous() |
| 1218 | param = params_dict[name] |
| 1219 | |
| 1220 | param.copy_(narrow_weight) |
| 1221 | loaded_params.add(name) |
| 1222 | continue |
| 1223 | elif ".w2_weight" in name: |
| 1224 | # Handle MLP down projection weights |
| 1225 | if use_ep: |
| 1226 | narrow_weight = weight[ep_rank_start:ep_rank_end, ...] |
| 1227 | else: |
| 1228 | narrow_weight = weight[:, tp_rank_start:tp_rank_end, :] |
| 1229 | narrow_weight = narrow_weight.permute(0, 2, 1).contiguous() |
| 1230 | param = params_dict[name] |
| 1231 | |
| 1232 | param.copy_(narrow_weight) |
| 1233 | loaded_params.add(name) |
| 1234 | continue |
| 1235 | elif ".w13_bias" in name: |
| 1236 | # Handle MLP gate and up projection biases |
| 1237 | # Extract gate and up projection bias parts |
| 1238 | if use_ep: |
| 1239 | narrow_weight = weight[ep_rank_start:ep_rank_end, ...] |
| 1240 | else: |
| 1241 | narrow_weight = weight[:, 2 * tp_rank_start : 2 * tp_rank_end] |
| 1242 | |
| 1243 | param = params_dict[name] |
| 1244 | param.copy_(narrow_weight) |
| 1245 | loaded_params.add(name) |
| 1246 | continue |
| 1247 | elif ".w2_bias" in name: |
| 1248 | # Handle MLP down projection bias |
| 1249 | if use_ep: |
| 1250 | weight = weight[ep_rank_start:ep_rank_end, ...] |
| 1251 | else: |
| 1252 | # (only load on rank 0 to avoid duplication) |
| 1253 | if tp_rank != 0: |
| 1254 | weight.zero_() |
| 1255 | param = params_dict[name] |
| 1256 | param.copy_(weight) |
| 1257 | loaded_params.add(name) |
| 1258 | continue |
| 1259 | elif "sinks" in name: |
| 1260 | # Handle attention sinks (distributed across ranks) |
| 1261 | param = params_dict[name] |
| 1262 | narrow_weight = weight.narrow(0, head_start, heads_per_rank) |
| 1263 | param.data.copy_(narrow_weight) |
| 1264 | loaded_params.add(name) |
| 1265 | continue |
| 1266 | for param_name, weight_name, shard_id in stacked_params_mapping: |
| 1267 | if weight_name not in name: |
| 1268 | continue |
| 1269 | name = name.replace(weight_name, param_name) |
| 1270 | param = params_dict[name] |
| 1271 | weight_loader = _get_weight_loader(param) |
| 1272 | if weight_loader == default_weight_loader: |
| 1273 | weight_loader(param, weight) |
| 1274 | else: |
| 1275 | weight_loader(param, weight, shard_id) |
| 1276 | break |
| 1277 | else: |
| 1278 | # Handle all other weights with potential renaming |
| 1279 | if name not in params_dict: |
| 1280 | continue |
| 1281 | param = params_dict[name] |
| 1282 | weight_loader = _get_weight_loader(param) |
| 1283 | weight_loader(param, weight) |
| 1284 | loaded_params.add(name) |
| 1285 | return loaded_params |
| 1286 | |
| 1287 | @staticmethod |
| 1288 | def _get_streamed_expert_info( |
| 1289 | name: str, |
| 1290 | params_dict: dict[str, torch.nn.Parameter], |
| 1291 | ) -> tuple[int, str, str] | None: |
| 1292 | """Parse ``...experts.<expert_id>.<fused_param>`` checkpoint keys.""" |
| 1293 | if ".mlp.experts." not in name: |
| 1294 | return None |
| 1295 | suffix = name.rsplit(".", 1)[-1] |
| 1296 | shard_id = _GPT_OSS_STREAMED_EXPERT_SUFFIX_TO_SHARD.get(suffix) |
| 1297 | if shard_id is None: |
| 1298 | return None |
| 1299 | |
| 1300 | prefix, expert_suffix = name.split(".mlp.experts.", maxsplit=1) |
| 1301 | expert_id_str, separator, param_suffix = expert_suffix.partition(".") |
| 1302 | if not separator or not expert_id_str.isdigit(): |
| 1303 | return None |
| 1304 | expert_id = int(expert_id_str) |
| 1305 | for base_layer_prefix in ("", "base_layer."): |
| 1306 | fused_name = ( |
| 1307 | f"{prefix}.mlp.experts.{base_layer_prefix}routed_experts.{param_suffix}" |
| 1308 | ) |
| 1309 | if fused_name in params_dict: |
| 1310 | return expert_id, fused_name, shard_id |
| 1311 | return None |
| 1312 | |
| 1313 | @classmethod |
| 1314 | def _try_load_streamed_expert( |
| 1315 | cls, |
| 1316 | name: str, |
| 1317 | loaded_weight: torch.Tensor, |
| 1318 | params_dict: dict[str, torch.nn.Parameter], |
| 1319 | loaded_params: set[str], |
| 1320 | ) -> bool: |
| 1321 | expert_info = cls._get_streamed_expert_info(name, params_dict) |
| 1322 | if expert_info is None: |
| 1323 | return False |
| 1324 | |
| 1325 | expert_id, fused_name, shard_id = expert_info |
| 1326 | param = params_dict[fused_name] |
| 1327 | weight_loader = typing.cast(Callable[..., bool], param.weight_loader) |
| 1328 | success = weight_loader( |
| 1329 | param, |
| 1330 | loaded_weight, |
| 1331 | weight_name=fused_name, |
| 1332 | shard_id=shard_id, |
| 1333 | expert_id=expert_id, |
| 1334 | return_success=True, |
| 1335 | ) |
| 1336 | if success: |
| 1337 | loaded_params.add(fused_name) |
| 1338 | return True |
| 1339 | |
| 1340 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 1341 | stacked_params_mapping = [ |
| 1342 | # (param_name, shard_name, shard_id) |
| 1343 | (".qkv_proj", ".q_proj", "q"), |
| 1344 | (".qkv_proj", ".k_proj", "k"), |
| 1345 | (".qkv_proj", ".v_proj", "v"), |
| 1346 | ] |
| 1347 | |
| 1348 | tp_rank = get_tensor_model_parallel_rank() |
| 1349 | tp_size = get_tensor_model_parallel_world_size() |
| 1350 | |
| 1351 | # Attention heads per rank |
| 1352 | heads_per_rank = self.config.num_attention_heads // tp_size |
| 1353 | head_start = tp_rank * heads_per_rank |
| 1354 | |
| 1355 | ep_size = get_ep_group().world_size |
| 1356 | ep_rank = get_ep_group().rank_in_group |
| 1357 | num_experts = self.config.num_local_experts |
| 1358 | experts_per_rank = num_experts // ep_size |
| 1359 | ep_rank_start = ep_rank * experts_per_rank |
| 1360 | ep_rank_end = (ep_rank + 1) * experts_per_rank |
| 1361 | |
| 1362 | quant_method = ( |
| 1363 | self.config.quantization_config["quant_method"] |
| 1364 | if hasattr(self.config, "quantization_config") |
| 1365 | else None |
| 1366 | ) |
| 1367 | # Normalize the checkpoint's quant_method to the internal name. |
| 1368 | # Note: there are three places where "mxfp4" -> "gpt_oss_mxfp4" |
| 1369 | # normalization occurs, each serving a different data path: |
| 1370 | # 1. GptOssMxfp4Config.override_quantization_method() — sets |
| 1371 | # ModelConfig.quantization (used to select the QuantizationConfig |
| 1372 | # class at model init time), reading from model_arch_config which |
| 1373 | # is a snapshot taken before verify_and_update_model_config runs. |
| 1374 | # 2. GptOssForCausalLMConfig.verify_and_update_model_config() — |
| 1375 | # patches hf_config.quantization_config in-place (a separate copy |
| 1376 | # of the dict from model_arch_config) for later hf_config lookups. |
| 1377 | # 3. Here — reads directly from self.config (the raw HF config) which |
| 1378 | # may still carry the original "mxfp4" string from the checkpoint. |
| 1379 | if quant_method == "mxfp4": |
| 1380 | quant_method = "gpt_oss_mxfp4" |
| 1381 | |
| 1382 | if quant_method == "gpt_oss_mxfp4": |
| 1383 | return self._load_weights_mxfp4( |
| 1384 | ep_rank_end, |
| 1385 | ep_rank_start, |
| 1386 | heads_per_rank, |
| 1387 | head_start, |
| 1388 | weights, |
| 1389 | stacked_params_mapping, |
| 1390 | ) |
| 1391 | elif quant_method == "quark": |
| 1392 | return self._load_weights_quark( |
| 1393 | ep_rank_end, |
| 1394 | ep_rank_start, |
| 1395 | heads_per_rank, |
| 1396 | head_start, |
| 1397 | weights, |
| 1398 | stacked_params_mapping, |
| 1399 | ) |
| 1400 | else: |
| 1401 | return self._load_weights_other( |
| 1402 | ep_rank_end, |
| 1403 | ep_rank_start, |
| 1404 | heads_per_rank, |
| 1405 | head_start, |
| 1406 | weights, |
| 1407 | stacked_params_mapping, |
| 1408 | ) |
| 1409 | |
| 1410 | |
| 1411 | class GptOssForCausalLM( |
| 1412 | nn.Module, SupportsPP, SupportsEagle, SupportsEagle3, SupportsLoRA |
| 1413 | ): |
| 1414 | is_3d_moe_weight: typing.ClassVar[bool] = True |
| 1415 | packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} |
| 1416 | |
| 1417 | hf_to_vllm_mapper = WeightsMapper( |
| 1418 | orig_to_new_substr={ |
| 1419 | ".self_attn.": ".attn.", |
| 1420 | }, |
| 1421 | orig_to_new_suffix={ |
| 1422 | ".embed_tokens.weight": ".embedding.weight", |
| 1423 | # MoE MXFP4 weights |
| 1424 | ".gate_up_proj_blocks": ".w13_weight", |
| 1425 | ".down_proj_blocks": ".w2_weight", |
| 1426 | ".gate_up_proj_scales": ".w13_weight_scale", |
| 1427 | ".down_proj_scales": ".w2_weight_scale", |
| 1428 | # MoE other weights |
| 1429 | ".gate_up_proj": ".w13_weight", |
| 1430 | ".down_proj": ".w2_weight", |
| 1431 | # MoE Bias |
| 1432 | ".gate_up_proj_bias": ".w13_bias", |
| 1433 | ".down_proj_bias": ".w2_bias", |
| 1434 | # For quark format |
| 1435 | ".gate_up_proj.weight": ".w13_weight", |
| 1436 | ".gate_up_proj.weight_scale": ".w13_weight_scale", |
| 1437 | ".gate_up_proj.bias": ".w13_bias", |
| 1438 | ".gate_up_proj.input_scale": ".w13_input_scale", |
| 1439 | ".down_proj.weight": ".w2_weight", |
| 1440 | ".down_proj.weight_scale": ".w2_weight_scale", |
| 1441 | ".down_proj.bias": ".w2_bias", |
| 1442 | ".down_proj.input_scale": ".w2_input_scale", |
| 1443 | }, |
| 1444 | ) |
| 1445 | |
| 1446 | def __init__( |
| 1447 | self, |
| 1448 | vllm_config: VllmConfig, |
| 1449 | prefix: str = "", |
| 1450 | ): |
| 1451 | super().__init__() |
| 1452 | self.vllm_config = vllm_config |
| 1453 | self.config = vllm_config.model_config.hf_config |
| 1454 | |
| 1455 | self.model = GptOssModel( |
| 1456 | vllm_config=vllm_config, |
| 1457 | prefix=maybe_prefix(prefix, "model"), |
| 1458 | ) |
| 1459 | self.lm_head = ParallelLMHead( |
| 1460 | self.config.vocab_size, |
| 1461 | self.config.hidden_size, |
| 1462 | prefix=maybe_prefix(prefix, "lm_head"), |
| 1463 | ) |
| 1464 | self.logits_processor = LogitsProcessor(self.config.vocab_size) |
| 1465 | self.make_empty_intermediate_tensors = ( |
| 1466 | self.model.make_empty_intermediate_tensors |
| 1467 | ) |
| 1468 | |
| 1469 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 1470 | return self.model.embed_input_ids(input_ids) |
| 1471 | |
| 1472 | def forward( |
| 1473 | self, |
| 1474 | input_ids: torch.Tensor | None, |
| 1475 | positions: torch.Tensor, |
| 1476 | intermediate_tensors: IntermediateTensors | None = None, |
| 1477 | inputs_embeds: torch.Tensor | None = None, |
| 1478 | ) -> torch.Tensor: |
| 1479 | return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) |
| 1480 | |
| 1481 | def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 1482 | logits = self.logits_processor(self.lm_head, hidden_states) |
| 1483 | return logits |
| 1484 | |
| 1485 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 1486 | loader = AutoWeightsLoader(self) |
| 1487 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |