# Copyright (c) 2026 LightSeek Foundation # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. from __future__ import annotations from collections.abc import Sequence from dataclasses import dataclass import torch from torch import nn from tokenspeed.runtime.layers.moe.schema import ExpertCheckpointSchema from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader @dataclass(frozen=True) class CheckpointPlanEntry: param_name: str checkpoint_weight_name: str shard_id: str def matches(self, checkpoint_name: str) -> bool: return self.checkpoint_weight_name in checkpoint_name def resolve_param_name(self, checkpoint_name: str) -> str: return checkpoint_name.replace(self.checkpoint_weight_name, self.param_name) @dataclass(frozen=True) class ExpertWeightPlanEntry(CheckpointPlanEntry): local_expert_id: int @dataclass(frozen=True) class FusedExpertWeightPlanEntry(CheckpointPlanEntry): split_dim: int | None = None split_chunks: int | None = None split_index: int | None = None class MoECheckpointLoadError(RuntimeError): pass def _build_default_expert_plan( schema: ExpertCheckpointSchema, *, num_experts: int, ep_rank: int, ep_size: int, ) -> list[ExpertWeightPlanEntry]: # Expert ownership is assumed to be a contiguous per-rank range here. # EPLB-aware remapping would need a different planning step. num_local_experts = num_experts // ep_size start_expert = num_local_experts * ep_rank expert_plan: list[ExpertWeightPlanEntry] = [] for local_expert_id in range(num_local_experts): expert_id = start_expert + local_expert_id expert_plan.extend( ( ExpertWeightPlanEntry( param_name="experts.w13_", checkpoint_weight_name=schema.make_expert_weight_name( expert_id, "gate_proj" ), shard_id="w1", local_expert_id=local_expert_id, ), ExpertWeightPlanEntry( param_name="experts.w13_", checkpoint_weight_name=schema.make_expert_weight_name( expert_id, "up_proj" ), shard_id="w3", local_expert_id=local_expert_id, ), ExpertWeightPlanEntry( param_name="experts.w2_", checkpoint_weight_name=schema.make_expert_weight_name( expert_id, "down_proj" ), shard_id="w2", local_expert_id=local_expert_id, ), ) ) return expert_plan def _build_global_expert_name_plan( schema: ExpertCheckpointSchema, *, num_experts: int, ) -> list[CheckpointPlanEntry]: expert_plan: list[CheckpointPlanEntry] = [] for expert_id in range(num_experts): expert_plan.extend( ( CheckpointPlanEntry( param_name="", checkpoint_weight_name=schema.make_expert_weight_name( expert_id, "gate_proj" ), shard_id="", ), CheckpointPlanEntry( param_name="", checkpoint_weight_name=schema.make_expert_weight_name( expert_id, "up_proj" ), shard_id="", ), CheckpointPlanEntry( param_name="", checkpoint_weight_name=schema.make_expert_weight_name( expert_id, "down_proj" ), shard_id="", ), ) ) return expert_plan def _build_default_fused_plan( schema: ExpertCheckpointSchema, *, fused_gate_up_as_w13: bool = False, include_bias: bool = False, ) -> list[FusedExpertWeightPlanEntry]: if fused_gate_up_as_w13: fused_plan = [ FusedExpertWeightPlanEntry( param_name="experts.w13_weight", checkpoint_weight_name=( f"experts.{schema.get_semantic_name('gate_up_fused')}" ), shard_id="w13", ), FusedExpertWeightPlanEntry( param_name="experts.w2_weight", checkpoint_weight_name=f"experts.{schema.get_semantic_name('down_proj')}", shard_id="w2", ), ] if include_bias: fused_plan.extend( ( FusedExpertWeightPlanEntry( param_name="experts.w13_weight_bias", checkpoint_weight_name=( f"experts.{schema.get_semantic_name('gate_up_bias')}" ), shard_id="w13", ), FusedExpertWeightPlanEntry( param_name="experts.w2_weight_bias", checkpoint_weight_name=( f"experts.{schema.get_semantic_name('down_bias')}" ), shard_id="w2", ), ) ) return fused_plan fused_plan = [ FusedExpertWeightPlanEntry( param_name="experts.w13_weight", checkpoint_weight_name=f"experts.{schema.get_semantic_name('gate_up_fused')}", shard_id="w1", split_dim=-2, split_chunks=2, split_index=0, ), FusedExpertWeightPlanEntry( param_name="experts.w13_weight", checkpoint_weight_name=f"experts.{schema.get_semantic_name('gate_up_fused')}", shard_id="w3", split_dim=-2, split_chunks=2, split_index=1, ), FusedExpertWeightPlanEntry( param_name="experts.w2_weight", checkpoint_weight_name=f"experts.{schema.get_semantic_name('down_proj')}", shard_id="w2", ), ] if include_bias: fused_plan.extend( ( FusedExpertWeightPlanEntry( param_name="experts.w13_weight_bias", checkpoint_weight_name=( f"experts.{schema.get_semantic_name('gate_up_bias')}" ), shard_id="w1", split_dim=-1, split_chunks=2, split_index=0, ), FusedExpertWeightPlanEntry( param_name="experts.w13_weight_bias", checkpoint_weight_name=( f"experts.{schema.get_semantic_name('gate_up_bias')}" ), shard_id="w3", split_dim=-1, split_chunks=2, split_index=1, ), FusedExpertWeightPlanEntry( param_name="experts.w2_weight_bias", checkpoint_weight_name=f"experts.{schema.get_semantic_name('down_bias')}", shard_id="w2", ), ) ) return fused_plan def _load_fused_expert_tensor( param, loaded_weight, *, shard_id: str, num_experts: int, ep_rank: int, ep_size: int, ) -> None: # Expert ownership is assumed to be a contiguous per-rank range here. # EPLB-aware remapping would need a different loading step. num_local_experts = num_experts // ep_size start_expert = num_local_experts * ep_rank end_expert = start_expert + num_local_experts weight_loader = param.weight_loader for expert_id in range(start_expert, end_expert): local_expert_id = expert_id - start_expert weight_loader( param, loaded_weight[expert_id], shard_id=shard_id, local_expert_id=local_expert_id, ) class MoECheckpointLoader: def __init__( self, *, params_dict: dict[str, nn.Parameter], expert_plan: Sequence[ExpertWeightPlanEntry] = (), global_expert_plan: Sequence[CheckpointPlanEntry] = (), fused_plan: Sequence[FusedExpertWeightPlanEntry] = (), num_experts: int | None = None, ep_rank: int = 0, ep_size: int = 1, fused_load_style: str = "per_expert", transpose_local_tensor_non_bias: bool = False, ) -> None: self._params_dict = params_dict self._expert_plan = tuple(expert_plan) self._global_expert_plan = tuple(global_expert_plan) self._fused_plan = tuple(fused_plan) self._num_experts = num_experts self._ep_rank = ep_rank self._ep_size = ep_size self._fused_load_style = fused_load_style self._transpose_local_tensor_non_bias = transpose_local_tensor_non_bias if self._fused_plan and self._num_experts is None: raise ValueError("num_experts is required when fused_plan is used") if fused_load_style not in {"per_expert", "local_tensor"}: raise ValueError(f"Unknown fused_load_style: {fused_load_style}") @staticmethod def _matches_plan(plan: Sequence[CheckpointPlanEntry], name: str) -> bool: return any(plan_entry.matches(name) for plan_entry in plan) def matches(self, name: str) -> bool: return self._matches_plan(self._fused_plan, name) or self._matches_plan( self._expert_plan, name ) def is_expert_checkpoint_weight(self, name: str) -> bool: """Return whether ``name`` belongs to this loader's MoE checkpoint schema. Args: name: Checkpoint tensor name after any model-specific remapping. Returns: ``True`` for local, non-local, or fused expert checkpoint tensors that this loader is responsible for; ``False`` for unrelated checkpoint tensors. """ return self._matches_plan(self._fused_plan, name) or self._matches_plan( self._global_expert_plan, name ) def _load_expert(self, name: str, loaded_weight: torch.Tensor) -> str | None: mapped_name: str | None = None for plan_entry in self._expert_plan: if not plan_entry.matches(name): continue mapped_name = plan_entry.resolve_param_name(name) param = self._params_dict.get(mapped_name) if param is None: continue param.weight_loader( param, loaded_weight, shard_id=plan_entry.shard_id, local_expert_id=plan_entry.local_expert_id, ) return mapped_name if mapped_name is not None: self._raise_unloaded_match(name, mapped_name) return None @staticmethod def _raise_unloaded_match(name: str, mapped_name: str | None) -> None: if mapped_name is None: raise MoECheckpointLoadError( f"Matched MoE checkpoint mapping for {name!r} but did not load any parameter" ) raise MoECheckpointLoadError( f"Matched MoE checkpoint mapping for {name!r} -> {mapped_name!r}, " "but the target parameter was not found or no tensor was loaded" ) @staticmethod def _raise_unmatched(name: str) -> None: raise MoECheckpointLoadError( f"{name!r} does not match any MoE checkpoint mapping" ) def _load_fused(self, name: str, loaded_weight: torch.Tensor) -> str | None: matched_entries = [ plan_entry for plan_entry in self._fused_plan if plan_entry.matches(name) ] if not matched_entries: return None selected_checkpoint_weight_name = max( (plan_entry.checkpoint_weight_name for plan_entry in matched_entries), key=len, ) loaded_any = False mapped_name: str | None = None for plan_entry in matched_entries: if plan_entry.checkpoint_weight_name != selected_checkpoint_weight_name: continue mapped_name = plan_entry.resolve_param_name(name) param = self._params_dict.get(mapped_name) if param is None: continue tensor_to_load = loaded_weight if plan_entry.split_dim is not None: tensor_to_load = loaded_weight.chunk( plan_entry.split_chunks, dim=plan_entry.split_dim )[plan_entry.split_index] if self._fused_load_style == "per_expert": _load_fused_expert_tensor( param, tensor_to_load, shard_id=plan_entry.shard_id, num_experts=self._num_experts, ep_rank=self._ep_rank, ep_size=self._ep_size, ) else: if self._transpose_local_tensor_non_bias and "bias" not in mapped_name: tensor_to_load = tensor_to_load.transpose(-2, -1) local_num_experts = param.shape[0] assert local_num_experts * self._ep_size == tensor_to_load.shape[0] local_experts = tensor_to_load[ local_num_experts * self._ep_rank : local_num_experts * (self._ep_rank + 1) ] if tensor_to_load.dtype == torch.float8_e5m2: default_weight_loader(param, local_experts.to(torch.bfloat16)) else: default_weight_loader(param, local_experts) loaded_any = True if not loaded_any: assert mapped_name is not None self._raise_unloaded_match(name, mapped_name) return mapped_name def load(self, name: str, loaded_weight: torch.Tensor) -> str: fused_mapped_name = self._load_fused(name, loaded_weight) if fused_mapped_name is not None: return fused_mapped_name expert_mapped_name = self._load_expert(name, loaded_weight) if expert_mapped_name is not None: return expert_mapped_name self._raise_unmatched(name) def build_moe_checkpoint_loader( *, params_dict: dict[str, nn.Parameter], expert_schema: ExpertCheckpointSchema | None = None, fused_schema: ExpertCheckpointSchema | None = None, num_experts: int | None = None, ep_rank: int = 0, ep_size: int = 1, fused_gate_up_as_w13: bool = False, include_bias: bool = False, fused_load_style: str = "per_expert", transpose_local_tensor_non_bias: bool = False, ) -> MoECheckpointLoader: expert_plan: Sequence[ExpertWeightPlanEntry] = () global_expert_plan: Sequence[CheckpointPlanEntry] = () if expert_schema is not None: if num_experts is None: raise ValueError("num_experts is required when expert_schema is used") expert_plan = _build_default_expert_plan( expert_schema, num_experts=num_experts, ep_rank=ep_rank, ep_size=ep_size, ) global_expert_plan = _build_global_expert_name_plan( expert_schema, num_experts=num_experts, ) fused_plan: Sequence[FusedExpertWeightPlanEntry] = () if fused_schema is not None: fused_plan = _build_default_fused_plan( fused_schema, fused_gate_up_as_w13=fused_gate_up_as_w13, include_bias=include_bias, ) return MoECheckpointLoader( params_dict=params_dict, expert_plan=expert_plan, global_expert_plan=global_expert_plan, fused_plan=fused_plan, num_experts=num_experts, ep_rank=ep_rank, ep_size=ep_size, fused_load_style=fused_load_style, transpose_local_tensor_non_bias=transpose_local_tensor_non_bias, )