from __future__ import annotations import json import re from pathlib import Path import torch import torch.nn as nn import torch.nn.functional as F from safetensors import safe_open from transformers import AutoModelForCausalLM, PreTrainedModel from .configuration_outlier_moe import OutlierMoEConfig def _parse_dtype(value): if value is None or value == "auto": return value if isinstance(value, torch.dtype): return value if isinstance(value, str): table = { "bfloat16": torch.bfloat16, "bf16": torch.bfloat16, "float16": torch.float16, "fp16": torch.float16, "float32": torch.float32, "fp32": torch.float32, } return table.get(value.lower(), value) return value def _load_alpha_map(model_dir: Path) -> dict[int, dict[int, float]]: path = model_dir / "alpha.json" raw = json.loads(path.read_text(encoding="utf-8")) out: dict[int, dict[int, float]] = {} for key, value in raw.items(): match = re.match(r"layer_(\d+)_expert_(\d+)", key) if not match: continue layer_idx = int(match.group(1)) expert_idx = int(match.group(2)) out.setdefault(layer_idx, {})[expert_idx] = float(value) return out def _load_router_map(model_dir: Path) -> dict[int, torch.Tensor]: path = model_dir / "router_state.safetensors" if not path.exists(): raise FileNotFoundError(f"Missing router state: {path}") out: dict[int, torch.Tensor] = {} with safe_open(str(path), framework="pt", device="cpu") as handle: for key in handle.keys(): match = re.match(r"layer_(\d+)_router_weight", key) if match: out[int(match.group(1))] = handle.get_tensor(key).float() if not out: raise RuntimeError(f"No router weights found in {path}") return out class CPUQuantizedExpert: def __init__(self, tensors: dict[str, torch.Tensor], alpha: float) -> None: self.gate_ternary = tensors["gate_ternary"].to(torch.int8).cpu() self.gate_scale = tensors["gate_scale"].to(torch.float16).cpu() self.up_ternary = tensors["up_ternary"].to(torch.int8).cpu() self.up_scale = tensors["up_scale"].to(torch.float16).cpu() self.down_ternary = tensors["down_ternary"].to(torch.int8).cpu() self.down_scale = tensors["down_scale"].to(torch.float16).cpu() self.alpha = float(alpha) def materialize(self, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: gate = self.gate_ternary.to(device=device, dtype=dtype) * self.gate_scale.to(device=device, dtype=dtype).unsqueeze(-1) up = self.up_ternary.to(device=device, dtype=dtype) * self.up_scale.to(device=device, dtype=dtype).unsqueeze(-1) down = self.down_ternary.to(device=device, dtype=dtype) * self.down_scale.to(device=device, dtype=dtype).unsqueeze(-1) return gate, up, down class EvalRoutedQuantizedMoE(nn.Module): def __init__(self, shared_mlp: nn.Module, experts: dict[int, CPUQuantizedExpert], router_weight: torch.Tensor, *, top_k: int) -> None: super().__init__() self.shared_mlp = shared_mlp self.experts = experts self.register_buffer("router_weight", router_weight.float().contiguous(), persistent=False) self.top_k = int(top_k) def forward(self, x: torch.Tensor) -> torch.Tensor: shared_out = self.shared_mlp(x) batch, seq_len, hidden = x.shape x_flat = x.reshape(-1, hidden) shared_flat = shared_out.reshape(-1, hidden) router_weight = self.router_weight.to(device=x.device, dtype=torch.float32) logits = F.linear(x_flat.float(), router_weight) vals, idx = torch.topk(logits, k=min(self.top_k, router_weight.shape[0]), dim=-1) weights = F.softmax(vals, dim=-1) mixed = shared_flat.float() target_dtype = x.dtype if x.dtype in (torch.float16, torch.bfloat16) else torch.float32 for expert_idx, expert in self.experts.items(): token_idx, choice_idx = torch.where(idx == expert_idx) if token_idx.numel() == 0: continue gate_w, up_w, down_w = expert.materialize(x.device, target_dtype) x_tok = x_flat[token_idx].to(dtype=target_dtype) gate = F.linear(x_tok, gate_w) up = F.linear(x_tok, up_w) out = F.linear(F.silu(gate) * up, down_w) delta = out.float() - shared_flat[token_idx].float() mixed[token_idx] += weights[token_idx, choice_idx].unsqueeze(-1) * expert.alpha * delta del gate_w, up_w, down_w, x_tok, gate, up, out, delta return mixed.to(dtype=shared_out.dtype).reshape(batch, seq_len, hidden) def _load_layer_experts(model_dir: Path, layer_idx: int, experts_per_layer: int, alpha_map: dict[int, dict[int, float]]) -> dict[int, CPUQuantizedExpert]: expert_dir = model_dir / "experts" layer_alphas = alpha_map.get(layer_idx, {}) experts: dict[int, CPUQuantizedExpert] = {} for expert_idx in range(experts_per_layer): path = expert_dir / f"layer_{layer_idx:02d}_expert_{expert_idx:02d}.safetensors" if not path.exists(): continue with safe_open(str(path), framework="pt", device="cpu") as handle: tensors = {key: handle.get_tensor(key) for key in handle.keys()} experts[expert_idx] = CPUQuantizedExpert(tensors, layer_alphas.get(expert_idx, 0.0)) return experts class OutlierMoEForCausalLM(PreTrainedModel): config_class = OutlierMoEConfig def __init__(self, config: OutlierMoEConfig) -> None: super().__init__(config) @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *model_args, config=None, **kwargs): model_dir = Path(pretrained_model_name_or_path) if config is None: config = OutlierMoEConfig.from_pretrained(model_dir) base_kwargs = {} for key in ("trust_remote_code", "device_map", "low_cpu_mem_usage", "attn_implementation"): if key in kwargs: base_kwargs[key] = kwargs.pop(key) torch_dtype = kwargs.pop("torch_dtype", None) if torch_dtype is None and "dtype" in kwargs: torch_dtype = kwargs.pop("dtype") if torch_dtype is not None: base_kwargs["torch_dtype"] = _parse_dtype(torch_dtype) model = AutoModelForCausalLM.from_pretrained( config.base_model_name_or_path, **base_kwargs, ) alpha_map = _load_alpha_map(model_dir) router_map = _load_router_map(model_dir) layers = list(getattr(config, "moe_layers", [])) experts_per_layer = int(getattr(config, "n_experts", 0)) top_k = int(getattr(config, "top_k", 2)) for layer_idx in layers: layer = model.model.layers[layer_idx] experts = _load_layer_experts(model_dir, layer_idx, experts_per_layer, alpha_map) router_weight = router_map[layer_idx] layer.mlp = EvalRoutedQuantizedMoE(layer.mlp, experts, router_weight, top_k=top_k) model.config = config return model