# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Inference-only Laguna model compatible with HuggingFace weights.""" import typing from collections.abc import Callable, Iterable from itertools import islice import torch import torch.nn.functional as F from torch import nn from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config from vllm.distributed import ( get_ep_group, get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( FusedMoE, fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, QKVParallelLinear, ReplicatedLinear, RowParallelLinear, ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, ) from vllm.model_executor.models.interfaces import ( EagleModelMixin, SupportsEagle3, SupportsLoRA, SupportsPP, ) from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, extract_layer_index, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, ) from vllm.sequence import IntermediateTensors logger = init_logger(__name__) class LagunaMLP(nn.Module): """Dense MLP for Laguna (used in mlp_only_layers).""" def __init__( self, hidden_size: int, intermediate_size: int, hidden_act: str, quant_config: QuantizationConfig | None = None, reduce_results: bool = True, prefix: str = "", ) -> None: super().__init__() # gate_proj and up_proj are kept as separate ColumnParallelLinear # rather than merged via MergedColumnParallelLinear. The merged form # requires per-partition NVFP4 global scales (weight_global_scale, # input_global_scale) to be packed into a length-2 PerTensorScaleParameter # and then collapsed via .max() in process_weights_after_loading; this # doesn't round-trip cleanly through Marlin's NVFP4 stacked-layer code # path. Splitting yields one global scale per Linear, exactly matching # the standard compressed-tensors per-Linear schema on disk. self.gate_proj = ColumnParallelLinear( hidden_size, intermediate_size, bias=False, quant_config=quant_config, prefix=f"{prefix}.gate_proj", ) self.up_proj = ColumnParallelLinear( hidden_size, intermediate_size, bias=False, quant_config=quant_config, prefix=f"{prefix}.up_proj", ) self.down_proj = RowParallelLinear( intermediate_size, hidden_size, bias=False, quant_config=quant_config, reduce_results=reduce_results, prefix=f"{prefix}.down_proj", ) if hidden_act != "silu": raise ValueError( f"Unsupported activation: {hidden_act}. Only silu is supported." ) def forward(self, x: torch.Tensor) -> torch.Tensor: gate, _ = self.gate_proj(x) up, _ = self.up_proj(x) x, _ = self.down_proj(F.silu(gate) * up) return x class LagunaMoE(nn.Module): """Sparse MoE block for Laguna with optional shared expert and sigmoid routing. Key differences from other MoE implementations: - Uses SIGMOID routing activation (not softmax) - Shared expert runs in parallel with routed experts (when enabled) - Matches HF reference: modular_laguna.py LagunaSparseMoeBlock """ def __init__( self, config, quant_config: QuantizationConfig | None = None, prefix: str = "", enable_eplb: bool = False, ): super().__init__() self.config = config self.num_experts = config.num_experts self.top_k = config.num_experts_per_tok self.tp_size = get_tensor_model_parallel_world_size() self.ep_group = get_ep_group().device_group self.ep_rank = self.ep_group.rank() self.ep_size = self.ep_group.size() self.n_routed_experts = config.num_experts self.n_shared_experts = 1 if config.shared_expert_intermediate_size > 0 else 0 self.routed_scaling_factor = float( getattr(config, "moe_routed_scaling_factor", 1.0) ) if self.tp_size > config.num_experts: raise ValueError( f"Tensor parallel size {self.tp_size} is greater than " f"the number of experts {config.num_experts}." ) # Load balancing settings. vllm_config = get_current_vllm_config() eplb_config = vllm_config.parallel_config.eplb_config self.enable_eplb = enable_eplb eplb_config.num_redundant_experts = ( eplb_config.num_redundant_experts if eplb_config.num_redundant_experts is not None else 0 ) self.n_redundant_experts = eplb_config.num_redundant_experts self.n_logical_experts = self.n_routed_experts self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts self.n_local_physical_experts = self.n_physical_experts // self.ep_size self.physical_expert_start = self.ep_rank * self.n_local_physical_experts self.physical_expert_end = ( self.physical_expert_start + self.n_local_physical_experts ) # Router gate self.gate = ReplicatedLinear( config.hidden_size, config.num_experts, bias=False, quant_config=None, prefix=f"{prefix}.gate", ) # Shared expert (optional) - passed to FusedMoE for overlap optimization self.shared_expert: LagunaMLP | None if config.shared_expert_intermediate_size > 0: self.shared_expert = LagunaMLP( hidden_size=config.hidden_size, intermediate_size=config.shared_expert_intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, reduce_results=False, # Reduce after shared+routed combine prefix=f"{prefix}.shared_expert", ) else: self.shared_expert = None # Auxiliary-loss-free load-balancing bias (arXiv:2408.15664). The # checkpoint stores one [num_experts] tensor per MoE layer at # `mlp.experts.e_score_correction_bias`; registering it as a Parameter # on the FusedMoE lets the weight loader pick it up and the router # add it during top-k selection. The fused top-k bias router requires # float32 regardless of model dtype. e_score_correction_bias = torch.nn.Parameter( torch.zeros(config.num_experts, dtype=torch.float32), requires_grad=False, ) # FusedMoE with SIGMOID routing. Passing `shared_experts=` lets the # layer overlap the shared-expert compute with the all2all dispatch. # `apply_routed_scale_to_output=True` makes FusedMoE handle the # routed_scaling_factor, shared+routed combine, and TP all-reduce # internally, so forward() just returns the final hidden states. self.experts = FusedMoE( shared_experts=self.shared_expert, num_experts=config.num_experts, top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, renormalize=config.norm_topk_prob, quant_config=quant_config, prefix=f"{prefix}.experts", scoring_func="sigmoid", use_grouped_topk=False, apply_router_weight_on_input=bool(config.moe_apply_router_weight_on_input), e_score_correction_bias=e_score_correction_bias, enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, routed_scaling_factor=self.routed_scaling_factor, apply_routed_scale_to_output=True, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: orig_shape = hidden_states.shape hidden_dim = hidden_states.shape[-1] hidden_states = hidden_states.view(-1, hidden_dim) router_logits, _ = self.gate(hidden_states) router_logits = router_logits.float() softcap = getattr(self.config, "moe_router_logit_softcapping", 0.0) or 0.0 if softcap > 0.0: router_logits = torch.tanh(router_logits / softcap) * softcap final_hidden_states = self.experts(hidden_states, router_logits) return final_hidden_states.view(orig_shape) class LagunaAttention(nn.Module): """Laguna attention with optional softplus output gating. Supports per-layer sliding window attention when ``config.layer_types`` is present. Layers whose type is ``"sliding_attention"`` use ``config.sliding_window``; all other layers (typically labelled ``"full_attention"``) use full attention. When ``layer_types`` is absent every layer defaults to full attention for backwards compatibility. """ def __init__( self, config, hidden_size: int, num_heads: int, num_kv_heads: int, max_position_embeddings: int = 131072, head_dim: int | None = None, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", attention_sink: bool = False, layer_idx: int | None = None, attention_prefix: str | None = None, ) -> None: super().__init__() if layer_idx is None: layer_idx = extract_layer_index(prefix) if attention_prefix is None: attention_prefix = prefix self.hidden_size = hidden_size tp_size = get_tensor_model_parallel_world_size() self.total_num_heads = num_heads assert self.total_num_heads % tp_size == 0 self.num_heads = self.total_num_heads // tp_size self.total_num_kv_heads = num_kv_heads if self.total_num_kv_heads >= tp_size: assert self.total_num_kv_heads % tp_size == 0 else: assert tp_size % self.total_num_kv_heads == 0 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) self.head_dim = head_dim or (hidden_size // self.total_num_heads) self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 self.max_position_embeddings = max_position_embeddings # Gating flag self.gating = config.gating # Per-layer sliding window (follows Gemma2/Cohere2 convention) layer_types = getattr(config, "layer_types", None) if layer_types is not None: is_sliding = layer_types[layer_idx] == "sliding_attention" self.sliding_window = config.sliding_window if is_sliding else None else: self.sliding_window = None # QKV projection (no bias for Laguna) self.qkv_proj = QKVParallelLinear( self.hidden_size, self.head_dim, self.total_num_heads, self.total_num_kv_heads, bias=config.qkv_bias, quant_config=quant_config, prefix=f"{prefix}.qkv_proj", ) # Output projection self.o_proj = RowParallelLinear( self.total_num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias, quant_config=quant_config, prefix=f"{prefix}.o_proj", ) # Gating projection (Laguna-specific, optional) # config.gating may be: # - True / "per-element": one gate per (head, head_dim) channel # - "per-head": one gate per head, broadcast across head_dim if self.gating: # v5 LagunaConfig uses ``gating=True`` for per-head; older configs # used ``"per-head"``. Accept both. ``"per-element"`` (or legacy # ``True``) means per-element gating with output size num_heads × # head_dim. gate_per_head = self.gating is True or self.gating == "per-head" g_out = ( self.total_num_heads if gate_per_head else self.total_num_heads * self.head_dim ) self.g_proj = ColumnParallelLinear( hidden_size, g_out, bias=False, quant_config=quant_config, prefix=f"{prefix}.g_proj", ) self.gate_per_head = gate_per_head else: self.g_proj = None self.gate_per_head = False # Attention sinks (learnable per-head bias for SWA layers) sinks = None if attention_sink: self.sink = torch.nn.Parameter( torch.empty(self.total_num_heads // tp_size, requires_grad=False) ) sinks = self.sink # Resolve rope params per-layer-type. ``config.rope_parameters`` is # either a flat dict (legacy) or a nested ``{layer_type: rope_dict}`` # (v5 Laguna-XS schema). The v5 form is unhashable as-is and would # crash `get_rope`'s cache lookup, so always pull out the layer's # sub-dict before forwarding. layer_type = ( layer_types[layer_idx] if layer_types is not None else "full_attention" ) is_sliding = layer_type == "sliding_attention" top_rope = getattr(config, "rope_parameters", None) or {} if any(isinstance(v, dict) for v in top_rope.values()): # Nested per-layer-type form. base_rope = top_rope.get(layer_type) or top_rope.get("full_attention") or {} else: base_rope = top_rope # Older flat-rope ckpts can carry a separate `swa_rope_parameters` # for SWA layers. Prefer it when present; otherwise the nested # rope dict above already supplies the correct sub-config. swa_rope = getattr(config, "swa_rope_parameters", None) if ( is_sliding and swa_rope is None and not any(isinstance(v, dict) for v in top_rope.values()) ): logger.warning_once( "Laguna config has sliding_attention layers but neither " "`swa_rope_parameters` nor a nested per-layer-type " "`rope_parameters` — SWA layers will reuse the global rope. " "If the checkpoint was trained with distinct SWA rope " "(theta / partial_rotary_factor), regenerate its HF config " "to include either form." ) rope_params = swa_rope if (is_sliding and swa_rope is not None) else base_rope # `partial_rotary_factor` may live on the top-level config (main attention) # or on the per-layer rope dict itself (e.g. SWA can differ). Inject the # top-level value into `rope_params` if the dict doesn't already set it. top_partial = getattr(config, "partial_rotary_factor", None) if top_partial is not None and "partial_rotary_factor" not in rope_params: rope_params = {**rope_params, "partial_rotary_factor": top_partial} # Rotary embeddings (YaRN) self.rotary_emb = get_rope( head_size=self.head_dim, max_position=max_position_embeddings, is_neox_style=True, rope_parameters=rope_params, ) self.attn = Attention( self.num_heads, self.head_dim, self.scaling, num_kv_heads=self.num_kv_heads, cache_config=cache_config, quant_config=quant_config, per_layer_sliding_window=self.sliding_window, prefix=f"{attention_prefix}.attn", sinks=sinks, ) # QK normalization (like Qwen3) self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim) q_by_head = self.q_norm(q_by_head) q = q_by_head.view(q.shape) k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim) k_by_head = self.k_norm(k_by_head) k = k_by_head.view(k.shape) q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v) # Apply gating if enabled (compute softplus in float32 for precision) if self.gating and self.g_proj is not None: gate, _ = self.g_proj(hidden_states) gate = F.softplus(gate.float()).type_as(attn_output) if self.gate_per_head: # gate: [..., num_heads]; broadcast across head_dim attn_shape = attn_output.shape attn_output = ( attn_output.view(*attn_shape[:-1], self.num_heads, self.head_dim) * gate.unsqueeze(-1) ).view(attn_shape) else: attn_output = attn_output * gate output, _ = self.o_proj(attn_output) return output class LagunaDecoderLayer(nn.Module): def __init__( self, config, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", enable_eplb: bool = False, layer_idx: int | None = None, attention_prefix: str | None = None, ) -> None: super().__init__() self.hidden_size = config.hidden_size if layer_idx is None: layer_idx = extract_layer_index(prefix) # Determine if this layer uses sliding window attention layer_types = getattr(config, "layer_types", None) is_sliding = ( layer_types is not None and layer_types[layer_idx] == "sliding_attention" ) # Enable attention sinks on SWA layers when configured attention_sink = is_sliding and getattr( config, "swa_attention_sink_enabled", False ) # Optional per-layer override of head count (Laguna-XS). per_layer_heads = getattr(config, "num_attention_heads_per_layer", None) layer_num_heads = ( per_layer_heads[layer_idx] if per_layer_heads is not None else config.num_attention_heads ) self.self_attn = LagunaAttention( config=config, hidden_size=self.hidden_size, num_heads=layer_num_heads, num_kv_heads=config.num_key_value_heads, max_position_embeddings=config.max_position_embeddings, head_dim=getattr(config, "head_dim", None), cache_config=cache_config, quant_config=quant_config, prefix=f"{prefix}.self_attn", attention_sink=attention_sink, layer_idx=layer_idx, attention_prefix=( f"{attention_prefix}.self_attn" if attention_prefix is not None else None ), ) # Check if this layer uses MoE or dense MLP (matches Qwen2/Qwen3 convention) mlp_only_layers = ( [] if not hasattr(config, "mlp_only_layers") else config.mlp_only_layers ) self.is_moe_layer = ( (layer_idx not in mlp_only_layers) and (config.num_experts > 0) and ((layer_idx + 1) % config.decoder_sparse_step == 0) ) if self.is_moe_layer: self.mlp = LagunaMoE( config=config, quant_config=quant_config, prefix=f"{prefix}.mlp", enable_eplb=enable_eplb, ) else: self.mlp = LagunaMLP( hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, prefix=f"{prefix}.mlp", ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, residual: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: # Self Attention if residual is None: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) else: hidden_states, residual = self.input_layernorm(hidden_states, residual) hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, ) # Fully Connected hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) hidden_states = self.mlp(hidden_states) return hidden_states, residual @support_torch_compile class LagunaModel(nn.Module, EagleModelMixin): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config cache_config = vllm_config.cache_config quant_config = vllm_config.quant_config enable_eplb = vllm_config.parallel_config.enable_eplb eplb_config = vllm_config.parallel_config.eplb_config self.num_redundant_experts = eplb_config.num_redundant_experts self.config = config self.quant_config = quant_config # Disable the model-level sliding-window fallback in Attention.__init__. # Laguna drives SWA per-layer via `layer_types`, passing # `per_layer_sliding_window=self.sliding_window` (None for global # layers). Without this, global layers whose `per_layer_sliding_window` # is None would pick up `cache_config.sliding_window` # (populated from `config.sliding_window`) as a fallback, silently # applying a 512-token window to full-attention layers. if cache_config is not None: cache_config.sliding_window = None self.vocab_size = config.vocab_size if get_pp_group().is_first_rank or ( config.tie_word_embeddings and get_pp_group().is_last_rank ): self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, quant_config=quant_config, prefix=f"{prefix}.embed_tokens", ) else: self.embed_tokens = PPMissingLayer() self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, lambda prefix: LagunaDecoderLayer( config=config, cache_config=cache_config, quant_config=quant_config, prefix=prefix, enable_eplb=enable_eplb, ), prefix=f"{prefix}.layers", ) if get_pp_group().is_last_rank: self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) else: self.norm = PPMissingLayer() self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, ) -> torch.Tensor | IntermediateTensors: if get_pp_group().is_first_rank: if inputs_embeds is not None: hidden_states = inputs_embeds else: hidden_states = self.embed_tokens(input_ids) residual = None else: assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] aux_hidden_states = self._maybe_add_hidden_state( [], self.start_layer, hidden_states, residual ) for layer_idx, layer in enumerate( islice(self.layers, self.start_layer, self.end_layer), start=self.start_layer, ): hidden_states, residual = layer(positions, hidden_states, residual) self._maybe_add_hidden_state( aux_hidden_states, layer_idx + 1, hidden_states, residual ) if not get_pp_group().is_last_rank: return IntermediateTensors( {"hidden_states": hidden_states, "residual": residual} ) hidden_states, _ = self.norm(hidden_states, residual) if len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states return hidden_states def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: """Get expert parameter mapping for weight loading. Returns mapping tuples of (param_name, weight_name, expert_id, shard_id) that handle both weights and quantization scales. """ return fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", num_experts=self.config.num_experts, num_redundant_experts=self.num_redundant_experts, ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), ("qkv_proj", "k_proj", "k"), ("qkv_proj", "v_proj", "v"), # gate_proj and up_proj are loaded as separate Linears (see # LagunaMLP) so no merge entry is needed here. ] # Suffixes to skip for GPTQ/modelopt models if param doesn't exist ignore_suffixes = ( ".bias", "_bias", ".k_scale", "_k_scale", ".v_scale", "_v_scale", ".weight_scale", "_weight_scale", ".input_scale", "_input_scale", ) params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() expert_params_mapping = self.get_expert_mapping() tp_rank = get_tensor_model_parallel_rank() for name, loaded_weight in weights: # Handle attention sinks (distributed across ranks). Derive the # per-rank slice from the parameter's own shape so per-layer # variations in head count are handled correctly. if "sink" in name: param = params_dict.get(name) if param is not None: layer_heads_per_rank = param.shape[0] layer_head_start = tp_rank * layer_heads_per_rank narrow_weight = loaded_weight.narrow( 0, layer_head_start, layer_heads_per_rank ) param.data.copy_(narrow_weight) loaded_params.add(name) continue # Handle stacked params (QKV, gate_up for # non-expert layers and shared_expert) for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue # Skip expert weights - handled below via expert_params_mapping if "mlp.experts" in name and "shared_expert" not in name: continue name = name.replace(weight_name, param_name) if name.endswith(ignore_suffixes) and name not in params_dict: continue if is_pp_missing_parameter(name, self): continue # Remap FP8 kv_scale names for backwards compatibility if name.endswith("scale"): name = maybe_remap_kv_scale_name(name, params_dict) if name is None: continue if name not in params_dict: continue param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) if weight_loader == default_weight_loader: weight_loader(param, loaded_weight) else: weight_loader(param, loaded_weight, shard_id) loaded_params.add(name) break else: # Try expert params mapping (handles weights + quantization scales) is_expert_weight = False for mapping in expert_params_mapping: param_name, weight_name, expert_id, shard_id = mapping if weight_name not in name: continue # Mark as expert weight so we skip regular loading below is_expert_weight = True # Create mapped name without modifying original name_mapped = name.replace(weight_name, param_name) if is_pp_missing_parameter(name_mapped, self): continue if ( name_mapped.endswith(ignore_suffixes) and name_mapped not in params_dict ): continue if name_mapped not in params_dict: continue param = params_dict[name_mapped] # Use return_success to handle expert parallelism correctly weight_loader = typing.cast( Callable[..., bool], param.weight_loader ) success = weight_loader( param, loaded_weight, name_mapped, shard_id=shard_id, expert_id=expert_id, return_success=True, ) if success: loaded_params.add(name_mapped) break else: # Expert weight not mapped to this rank - skip if is_expert_weight: continue # Remap kv_scale names before the ignore_suffixes filter: # the suffix list includes .k_scale/.v_scale, so filtering # first drops the checkpoint key before remap can rewrite # it to the .attn.* name that exists in params_dict. name = maybe_remap_kv_scale_name(name, params_dict) if name is None: continue if name.endswith(ignore_suffixes) and name not in params_dict: continue if is_pp_missing_parameter(name, self): continue if name not in params_dict: continue param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params class LagunaForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SupportsEagle3): fall_back_to_pt_during_load = False packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config self.config = config self.quant_config = quant_config self.model = LagunaModel( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) if get_pp_group().is_last_rank: self.lm_head = ParallelLMHead( config.vocab_size, config.hidden_size, quant_config=quant_config, prefix=maybe_prefix(prefix, "lm_head"), ) if self.config.tie_word_embeddings: self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) else: self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, ) -> torch.Tensor | IntermediateTensors: hidden_states = self.model( input_ids, positions, intermediate_tensors, inputs_embeds ) return hidden_states def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: logits = self.logits_processor(self.lm_head, hidden_states) return logits def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping() def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader( self, skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), ) return loader.load_weights(weights)