680 lines
24 KiB
Python
680 lines
24 KiB
Python
"""
|
|
LFM2-MoE (Liquid Foundation Model 2 - Mixture of Experts) implementation for SGLang.
|
|
|
|
This is a hybrid architecture with attention, ShortConv, and MoE layers:
|
|
- Attention layers use standard KV cache (RadixAttention)
|
|
- Conv layers use MambaPool for state caching (via HybridReqToTokenPool)
|
|
- First `num_dense_layers` use dense MLP, rest use MoE with sigmoid routing
|
|
|
|
Key MoE characteristics:
|
|
- Sigmoid routing (not softmax) - auxiliary-loss-free style
|
|
- Expert bias (fp32) affects selection but not weighting
|
|
- Post-hoc normalization of top-k weights
|
|
"""
|
|
|
|
from typing import Iterable, Optional, Set, Tuple
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
from sglang.srt.configs.lfm2_moe import Lfm2MoeConfig
|
|
from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
|
|
from sglang.srt.layers.activation import SiluAndMul
|
|
from sglang.srt.layers.attention.mamba.causal_conv1d import (
|
|
causal_conv1d_fn,
|
|
causal_conv1d_update,
|
|
)
|
|
from sglang.srt.layers.layernorm import RMSNorm
|
|
from sglang.srt.layers.linear import (
|
|
MergedColumnParallelLinear,
|
|
QKVParallelLinear,
|
|
ReplicatedLinear,
|
|
RowParallelLinear,
|
|
)
|
|
from sglang.srt.layers.logits_processor import LogitsProcessor
|
|
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
|
|
from sglang.srt.layers.moe.topk import TopK
|
|
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
|
from sglang.srt.layers.radix_attention import RadixAttention
|
|
from sglang.srt.layers.rotary_embedding import get_rope
|
|
from sglang.srt.layers.vocab_parallel_embedding import (
|
|
ParallelLMHead,
|
|
VocabParallelEmbedding,
|
|
)
|
|
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
|
from sglang.srt.model_loader.weight_utils import (
|
|
default_weight_loader,
|
|
sharded_weight_loader,
|
|
)
|
|
from sglang.srt.utils import add_prefix, make_layers, set_weight_attrs
|
|
|
|
|
|
class Lfm2MoeMLP(nn.Module):
|
|
"""Dense MLP for first N layers (before MoE kicks in)."""
|
|
|
|
def __init__(
|
|
self,
|
|
config: Lfm2MoeConfig,
|
|
quant_config: Optional[QuantizationConfig] = None,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
# Use MergedColumnParallelLinear for w1/w3 (gate/up projections)
|
|
self.gate_up_proj = MergedColumnParallelLinear(
|
|
config.hidden_size,
|
|
[config.intermediate_size] * 2,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("gate_up_proj", prefix),
|
|
)
|
|
self.down_proj = RowParallelLinear(
|
|
config.intermediate_size,
|
|
config.hidden_size,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("down_proj", prefix),
|
|
)
|
|
self.act_fn = SiluAndMul()
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
gate_up, _ = self.gate_up_proj(x)
|
|
x = self.act_fn(gate_up)
|
|
out, _ = self.down_proj(x)
|
|
return out
|
|
|
|
|
|
class Lfm2MoeSparseMoeBlock(nn.Module):
|
|
"""
|
|
Sparse MoE block with sigmoid routing using optimized FusedMoE.
|
|
|
|
Key features:
|
|
- Sigmoid scoring (not softmax) - auxiliary-loss-free style
|
|
- Expert bias (fp32) for load balancing
|
|
- Bias affects selection only, not weighting
|
|
- Uses FusedMoE for efficient batched expert computation
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
config: Lfm2MoeConfig,
|
|
layer_idx: int,
|
|
quant_config: Optional[QuantizationConfig] = None,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
self.tp_size = get_tensor_model_parallel_world_size()
|
|
self.routed_scaling_factor = config.routed_scaling_factor
|
|
|
|
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}."
|
|
)
|
|
|
|
# Gate (router) - outputs logits for each expert
|
|
self.gate = ReplicatedLinear(
|
|
config.hidden_size,
|
|
config.num_experts,
|
|
bias=False,
|
|
quant_config=None,
|
|
prefix=add_prefix("gate", prefix),
|
|
)
|
|
|
|
# Expert bias (fp32) - affects selection but not weighting
|
|
if config.use_expert_bias:
|
|
self.expert_bias = nn.Parameter(
|
|
torch.zeros(config.num_experts, dtype=torch.float32)
|
|
)
|
|
else:
|
|
self.register_parameter("expert_bias", None)
|
|
|
|
# TopK selector with sigmoid scoring
|
|
self.topk = TopK(
|
|
top_k=config.num_experts_per_tok,
|
|
layer_id=layer_idx,
|
|
renormalize=config.norm_topk_prob,
|
|
scoring_func="sigmoid",
|
|
correction_bias=self.expert_bias if config.use_expert_bias else None,
|
|
)
|
|
|
|
# FusedMoE for efficient batched expert computation
|
|
# Note: We intentionally do NOT pass routed_scaling_factor to FusedMoE.
|
|
# While FusedMoE supports it, passing it there increases numerical
|
|
# differences vs HuggingFace (likely due to different code paths in the
|
|
# Triton runner when scaling_factor != None). We apply it manually below.
|
|
self.experts = FusedMoE(
|
|
num_experts=config.num_experts,
|
|
top_k=config.num_experts_per_tok,
|
|
hidden_size=config.hidden_size,
|
|
intermediate_size=config.moe_intermediate_size,
|
|
layer_id=layer_idx,
|
|
reduce_results=True,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("experts", prefix),
|
|
)
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
"""Optimized expert forward pass using FusedMoE."""
|
|
# Get router logits
|
|
router_logits, _ = self.gate(hidden_states)
|
|
|
|
# Select top-k experts with sigmoid scoring
|
|
topk_output = self.topk(hidden_states, router_logits)
|
|
|
|
# Run fused expert computation
|
|
final_hidden_states = self.experts(hidden_states, topk_output)
|
|
|
|
# Apply routed scaling factor (see __init__ comment for why not in FusedMoE)
|
|
return final_hidden_states * self.routed_scaling_factor
|
|
|
|
|
|
class Lfm2MoeAttention(nn.Module):
|
|
"""Grouped-query attention with RoPE and Q/K layernorm."""
|
|
|
|
def __init__(
|
|
self,
|
|
config: Lfm2MoeConfig,
|
|
layer_id: int,
|
|
quant_config: Optional[QuantizationConfig] = None,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
self.hidden_size = config.hidden_size
|
|
self.total_num_heads = config.num_attention_heads
|
|
self.total_num_kv_heads = config.num_key_value_heads
|
|
self.head_dim = self.hidden_size // self.total_num_heads
|
|
self.scaling = self.head_dim**-0.5
|
|
|
|
rope_parameters = getattr(config, "rope_parameters", None)
|
|
if rope_parameters is not None and "rope_theta" in rope_parameters:
|
|
rope_theta = rope_parameters["rope_theta"]
|
|
else:
|
|
rope_theta = getattr(config, "rope_theta", 1000000.0)
|
|
|
|
self.rotary_emb = get_rope(
|
|
head_size=self.head_dim,
|
|
rotary_dim=self.head_dim,
|
|
max_position=getattr(config, "max_position_embeddings", 128000),
|
|
rope_scaling=getattr(config, "rope_scaling", None),
|
|
base=rope_theta,
|
|
is_neox_style=True,
|
|
dtype=torch.get_default_dtype(),
|
|
)
|
|
|
|
self.qkv_proj = QKVParallelLinear(
|
|
self.hidden_size,
|
|
self.head_dim,
|
|
self.total_num_heads,
|
|
self.total_num_kv_heads,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("qkv_proj", prefix),
|
|
)
|
|
self.out_proj = RowParallelLinear(
|
|
self.total_num_heads * self.head_dim,
|
|
self.hidden_size,
|
|
bias=False,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("out_proj", prefix),
|
|
)
|
|
|
|
self.q_layernorm = RMSNorm(self.head_dim, eps=config.norm_eps)
|
|
self.k_layernorm = RMSNorm(self.head_dim, eps=config.norm_eps)
|
|
|
|
self.num_local_q_heads = self.qkv_proj.num_heads
|
|
self.num_local_kv_heads = self.qkv_proj.num_kv_heads
|
|
|
|
self.attn = RadixAttention(
|
|
num_heads=self.num_local_q_heads,
|
|
head_dim=self.head_dim,
|
|
scaling=self.scaling,
|
|
num_kv_heads=self.num_local_kv_heads,
|
|
layer_id=layer_id,
|
|
prefix=add_prefix("attn", prefix),
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
positions: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
forward_batch: ForwardBatch,
|
|
) -> torch.Tensor:
|
|
T = hidden_states.shape[0]
|
|
qkv, _ = self.qkv_proj(hidden_states)
|
|
|
|
q_size = self.num_local_q_heads * self.head_dim
|
|
kv_size = self.num_local_kv_heads * self.head_dim
|
|
q, k, v = torch.split(qkv, [q_size, kv_size, kv_size], dim=-1)
|
|
|
|
q = q.reshape(T, self.num_local_q_heads, self.head_dim)
|
|
k = k.reshape(T, self.num_local_kv_heads, self.head_dim)
|
|
|
|
q = self.q_layernorm(q.reshape(-1, self.head_dim)).reshape(
|
|
T, self.num_local_q_heads, self.head_dim
|
|
)
|
|
k = self.k_layernorm(k.reshape(-1, self.head_dim)).reshape(
|
|
T, self.num_local_kv_heads, self.head_dim
|
|
)
|
|
|
|
q, k = self.rotary_emb(positions, q, k)
|
|
|
|
attn_out = self.attn(q.reshape(T, -1), k.reshape(T, -1), v, forward_batch)
|
|
out, _ = self.out_proj(attn_out)
|
|
return out
|
|
|
|
|
|
class Lfm2MoeShortConv(nn.Module):
|
|
"""
|
|
Gated short convolution layer using optimized causal_conv1d kernels.
|
|
|
|
Architecture: in_proj -> split(B, C, x) -> Bx -> conv1d -> C*conv_out -> out_proj
|
|
- Supports tensor parallelism: hidden dimension is sharded across TP ranks
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
config: Lfm2MoeConfig,
|
|
layer_idx: int,
|
|
quant_config: Optional[QuantizationConfig] = None,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
self.layer_idx = layer_idx
|
|
self.conv_kernel = int(config.conv_L_cache)
|
|
self.use_bias = bool(config.conv_bias)
|
|
self.hidden_size = config.hidden_size
|
|
|
|
# Get tensor parallel size for sharding
|
|
self.tp_size = get_tensor_model_parallel_world_size()
|
|
self.hidden_size_per_partition = self.hidden_size // self.tp_size
|
|
|
|
# Use MergedColumnParallelLinear so each output (B, C, x) is sharded separately
|
|
self.in_proj = MergedColumnParallelLinear(
|
|
config.hidden_size,
|
|
[config.hidden_size] * 3, # B, C, x each get hidden_size
|
|
bias=self.use_bias,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.in_proj",
|
|
)
|
|
self.out_proj = RowParallelLinear(
|
|
config.hidden_size,
|
|
config.hidden_size,
|
|
bias=self.use_bias,
|
|
input_is_parallel=True,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.out_proj",
|
|
)
|
|
|
|
# Conv weights sharded along hidden dimension: (hidden_size/tp, kernel_size)
|
|
self.conv_weight = nn.Parameter(
|
|
torch.empty(self.hidden_size_per_partition, self.conv_kernel)
|
|
)
|
|
set_weight_attrs(self.conv_weight, {"weight_loader": sharded_weight_loader(0)})
|
|
if self.use_bias:
|
|
self.conv_bias = nn.Parameter(torch.empty(self.hidden_size_per_partition))
|
|
set_weight_attrs(
|
|
self.conv_bias, {"weight_loader": sharded_weight_loader(0)}
|
|
)
|
|
else:
|
|
self.register_parameter("conv_bias", None)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
forward_batch: ForwardBatch,
|
|
) -> torch.Tensor:
|
|
if forward_batch.forward_mode.is_idle():
|
|
return hidden_states
|
|
|
|
layer_cache = forward_batch.req_to_token_pool.mamba2_layer_cache(self.layer_idx)
|
|
conv_state = layer_cache.conv[0]
|
|
req_pool_indices = forward_batch.req_pool_indices
|
|
|
|
proj, _ = self.in_proj(hidden_states)
|
|
B_gate, C_gate, x = proj.chunk(3, dim=-1)
|
|
Bx = B_gate * x
|
|
|
|
if forward_batch.forward_mode.is_decode():
|
|
conv_out = causal_conv1d_update(
|
|
Bx,
|
|
conv_state,
|
|
self.conv_weight,
|
|
self.conv_bias,
|
|
activation=None,
|
|
conv_state_indices=req_pool_indices.to(torch.int32),
|
|
)
|
|
else:
|
|
T = hidden_states.shape[0]
|
|
Bx_t = Bx.transpose(0, 1).contiguous()
|
|
|
|
# Build query_start_loc for variable-length sequences
|
|
# causal_conv1d_fn expects [start0, start1, ..., startN, T]
|
|
extend_start_loc = forward_batch.extend_start_loc
|
|
if extend_start_loc is not None and len(extend_start_loc) > 1:
|
|
# Multiple sequences: append T to extend_start_loc
|
|
# Allocate and fill to avoid torch.cat overhead
|
|
query_start_loc = extend_start_loc.new_empty(len(extend_start_loc) + 1)
|
|
query_start_loc[:-1] = extend_start_loc
|
|
query_start_loc[-1] = T
|
|
cache_indices = req_pool_indices.to(torch.int32)
|
|
else:
|
|
# Single sequence: [0, T]
|
|
query_start_loc = hidden_states.new_tensor([0, T], dtype=torch.int32)
|
|
cache_indices = req_pool_indices[:1].to(torch.int32)
|
|
|
|
conv_out = causal_conv1d_fn(
|
|
Bx_t,
|
|
self.conv_weight,
|
|
self.conv_bias,
|
|
query_start_loc=query_start_loc,
|
|
cache_indices=cache_indices,
|
|
has_initial_state=None,
|
|
conv_states=conv_state,
|
|
activation=None,
|
|
).transpose(0, 1)
|
|
|
|
output, _ = self.out_proj(C_gate * conv_out)
|
|
return output
|
|
|
|
|
|
class Lfm2MoeDecoderLayer(nn.Module):
|
|
"""
|
|
Decoder layer with attention/conv and dense MLP or MoE.
|
|
|
|
- Layers 0 to num_dense_layers-1: use Lfm2MoeMLP (dense)
|
|
- Layers num_dense_layers+: use Lfm2MoeSparseMoeBlock (MoE)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
config: Lfm2MoeConfig,
|
|
layer_id: int,
|
|
quant_config: Optional[QuantizationConfig] = None,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
self.layer_type = config.layer_types[layer_id]
|
|
self.is_attention_layer = self.layer_type == "full_attention"
|
|
|
|
self.operator_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
|
|
self.ffn_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
|
|
|
|
# Attention or Conv
|
|
if self.is_attention_layer:
|
|
self.self_attn = Lfm2MoeAttention(
|
|
config=config,
|
|
layer_id=layer_id,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("self_attn", prefix),
|
|
)
|
|
else:
|
|
self.conv = Lfm2MoeShortConv(
|
|
config=config,
|
|
layer_idx=layer_id,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("conv", prefix),
|
|
)
|
|
|
|
# Dense MLP or MoE
|
|
if layer_id < config.num_dense_layers:
|
|
self.feed_forward = Lfm2MoeMLP(
|
|
config=config,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("feed_forward", prefix),
|
|
)
|
|
else:
|
|
self.feed_forward = Lfm2MoeSparseMoeBlock(
|
|
config=config,
|
|
layer_idx=layer_id,
|
|
quant_config=quant_config,
|
|
prefix=add_prefix("feed_forward", prefix),
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
layer_id: int,
|
|
positions: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
residual: Optional[torch.Tensor],
|
|
forward_batch: ForwardBatch,
|
|
**kwargs,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
if not forward_batch.forward_mode.is_idle():
|
|
residual = hidden_states
|
|
normed = self.operator_norm(hidden_states)
|
|
|
|
if self.is_attention_layer:
|
|
hidden_states = self.self_attn(positions, normed, forward_batch)
|
|
else:
|
|
hidden_states = self.conv(normed, forward_batch)
|
|
|
|
hidden_states = hidden_states + residual
|
|
hidden_states = hidden_states + self.feed_forward(
|
|
self.ffn_norm(hidden_states)
|
|
)
|
|
|
|
return hidden_states, residual
|
|
|
|
|
|
class Lfm2MoeModel(nn.Module):
|
|
def __init__(
|
|
self,
|
|
config: Lfm2MoeConfig,
|
|
quant_config: Optional[QuantizationConfig] = None,
|
|
prefix: str = "",
|
|
):
|
|
super().__init__()
|
|
self.config = config
|
|
|
|
self.embed_tokens = VocabParallelEmbedding(
|
|
config.vocab_size,
|
|
config.hidden_size,
|
|
org_num_embeddings=config.vocab_size,
|
|
prefix=add_prefix("embed_tokens", prefix),
|
|
)
|
|
|
|
# Count attention layers for KV cache sizing
|
|
self.num_attention_layers = sum(
|
|
1 for lt in config.layer_types if lt == "full_attention"
|
|
)
|
|
|
|
def get_layer(idx: int, prefix: str, **kwargs):
|
|
return Lfm2MoeDecoderLayer(
|
|
config=config,
|
|
layer_id=idx,
|
|
quant_config=quant_config,
|
|
prefix=prefix,
|
|
)
|
|
|
|
self.layers = make_layers(
|
|
config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
|
|
)
|
|
self.embedding_norm = RMSNorm(config.hidden_size, eps=config.norm_eps)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
forward_batch: ForwardBatch,
|
|
inputs_embeds: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
hidden_states = (
|
|
inputs_embeds if inputs_embeds is not None else self.embed_tokens(input_ids)
|
|
)
|
|
|
|
residual = None
|
|
for i in range(len(self.layers)):
|
|
hidden_states, residual = self.layers[i](
|
|
layer_id=i,
|
|
positions=positions,
|
|
hidden_states=hidden_states,
|
|
residual=residual,
|
|
forward_batch=forward_batch,
|
|
)
|
|
|
|
return self.embedding_norm(hidden_states)
|
|
|
|
|
|
class Lfm2MoeForCausalLM(nn.Module):
|
|
"""LFM2-MoE for causal language modeling."""
|
|
|
|
fall_back_to_pt_during_load = False
|
|
|
|
def __init__(
|
|
self,
|
|
config: Lfm2MoeConfig,
|
|
quant_config: Optional[QuantizationConfig] = None,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
self.config = config
|
|
self.pp_group = get_pp_group()
|
|
assert self.pp_group.is_first_rank and self.pp_group.is_last_rank
|
|
|
|
self.quant_config = quant_config
|
|
self.model = Lfm2MoeModel(
|
|
config, quant_config, prefix=add_prefix("model", prefix)
|
|
)
|
|
self.lm_head = ParallelLMHead(
|
|
config.vocab_size,
|
|
config.hidden_size,
|
|
quant_config=quant_config,
|
|
org_num_embeddings=config.vocab_size,
|
|
prefix=add_prefix("lm_head", prefix),
|
|
)
|
|
self.logits_processor = LogitsProcessor(config)
|
|
self.num_attention_layers = self.model.num_attention_layers
|
|
|
|
def get_num_kv_cache_layers(self) -> int:
|
|
return self.num_attention_layers
|
|
|
|
@torch.no_grad()
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
forward_batch: ForwardBatch,
|
|
inputs_embeds: Optional[torch.Tensor] = None,
|
|
**kwargs,
|
|
):
|
|
hidden_states = self.model(input_ids, positions, forward_batch, inputs_embeds)
|
|
return self.logits_processor(
|
|
input_ids, hidden_states, self.lm_head, forward_batch
|
|
)
|
|
|
|
def load_weights(
|
|
self, weights: Iterable[Tuple[str, torch.Tensor]], is_mtp: bool = False
|
|
) -> Set[str]:
|
|
"""Load weights with FusedMoE expert format."""
|
|
stacked_params_mapping = [
|
|
# (param_name, weight_name, shard_id)
|
|
("qkv_proj", "q_proj", "q"),
|
|
("qkv_proj", "k_proj", "k"),
|
|
("qkv_proj", "v_proj", "v"),
|
|
# Dense MLP w1/w3 -> gate_up_proj
|
|
("gate_up_proj", "w1", 0),
|
|
("gate_up_proj", "w3", 1),
|
|
]
|
|
|
|
# FusedMoE expert params mapping
|
|
# HF format: experts.{expert_id}.w{1,2,3}.weight
|
|
# FusedMoE format: experts.w13_weight, experts.w2_weight
|
|
expert_params_mapping = FusedMoE.make_expert_params_mapping(
|
|
ckpt_gate_proj_name="w1",
|
|
ckpt_down_proj_name="w2",
|
|
ckpt_up_proj_name="w3",
|
|
num_experts=self.config.num_experts,
|
|
)
|
|
|
|
params_dict = dict(self.named_parameters())
|
|
loaded_params: Set[str] = set()
|
|
embed_tokens_weight = None
|
|
|
|
for name, loaded_weight in weights:
|
|
if "rotary_emb.inv_freq" in name:
|
|
continue
|
|
|
|
if "embed_tokens.weight" in name:
|
|
embed_tokens_weight = loaded_weight
|
|
|
|
# Handle conv weight/bias naming: HF uses conv.conv, we use conv_weight/conv_bias
|
|
if ".conv.conv.weight" in name:
|
|
name = name.replace(".conv.conv.weight", ".conv.conv_weight")
|
|
loaded_weight = loaded_weight.squeeze(1) # (D, 1, K) -> (D, K)
|
|
if ".conv.conv.bias" in name:
|
|
name = name.replace(".conv.conv.bias", ".conv.conv_bias")
|
|
|
|
# Handle dense MLP w2 -> down_proj
|
|
if "feed_forward.w2" in name and "experts" not in name:
|
|
name = name.replace("feed_forward.w2", "feed_forward.down_proj")
|
|
|
|
# Handle stacked params (QKV, dense MLP gate_up)
|
|
for param_name, weight_name, shard_id in stacked_params_mapping:
|
|
if weight_name not in name:
|
|
continue
|
|
# Skip expert weights (handled below)
|
|
if "experts" in name:
|
|
continue
|
|
name = name.replace(weight_name, param_name)
|
|
if name.endswith(".bias") and name not in params_dict:
|
|
break
|
|
if name not in params_dict:
|
|
break
|
|
param = params_dict[name]
|
|
weight_loader = getattr(param, "weight_loader")
|
|
weight_loader(param, loaded_weight, shard_id)
|
|
loaded_params.add(name)
|
|
break
|
|
else:
|
|
# Handle MoE expert weights using FusedMoE format
|
|
# HF format: model.layers.X.feed_forward.experts.Y.wZ.weight
|
|
# FusedMoE format: model.layers.X.feed_forward.experts.w13_weight/w2_weight
|
|
for (
|
|
param_name,
|
|
weight_name,
|
|
expert_id,
|
|
shard_id,
|
|
) in expert_params_mapping:
|
|
if weight_name not in name:
|
|
continue
|
|
# Build our parameter name
|
|
name = name.replace(weight_name, param_name)
|
|
if name not in params_dict:
|
|
continue
|
|
param = params_dict[name]
|
|
weight_loader = param.weight_loader
|
|
weight_loader(
|
|
param,
|
|
loaded_weight,
|
|
name,
|
|
shard_id=shard_id,
|
|
expert_id=expert_id,
|
|
)
|
|
loaded_params.add(name)
|
|
break
|
|
else:
|
|
# Handle regular weights
|
|
if name.endswith(".bias") and name not in params_dict:
|
|
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)
|
|
|
|
# Handle tied lm_head weight
|
|
if "lm_head.weight" not in loaded_params and "lm_head.weight" in params_dict:
|
|
if embed_tokens_weight is not None:
|
|
param = params_dict["lm_head.weight"]
|
|
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
|
weight_loader(param, embed_tokens_weight)
|
|
loaded_params.add("lm_head.weight")
|
|
|
|
return loaded_params
|
|
|
|
|
|
EntryClass = [Lfm2MoeForCausalLM]
|