96 lines
3.5 KiB
Python
96 lines
3.5 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""
|
|||
|
|
ixformer.contrib.vllm.layers —— 纯 PyTorch 回退（R7 兜底）
|
|||
|
|
|
|||
|
|
存在原因（实证）：Iluvatar bi-100 引擎镜像里没有 `ixformer.contrib.vllm.layers`，
|
|||
|
|
MoE 模型实现里
|
|||
|
|
from ixformer.contrib.vllm.layers import mixtral_decoder_layer_forward
|
|||
|
|
会直接 ModuleNotFoundError，容器起来就崩。缺的是**引擎镜像的依赖**，不是模型的问题。
|
|||
|
|
|
|||
|
|
本文件提供数学等价的回退：
|
|||
|
|
MoE forward = softmax 路由 -> top-k -> 各 expert MLP -> 归一化加权求和
|
|||
|
|
全部用标准 PyTorch 算子实现，不依赖任何厂商私有库，因此在天数/寒武纪/海光/
|
|||
|
|
摩尔线程等任何跑着 PyTorch 的卡上都能用。
|
|||
|
|
|
|||
|
|
定位说明：这是"让它能跑"的兜底，不是"跑得快"的优化。追求吞吐仍应用厂商原生算子。
|
|||
|
|
只有真缺模块时才会被 import 到（经 entrypoint 注入 PYTHONPATH=/opt/shims）。
|
|||
|
|
"""
|
|||
|
|
import torch
|
|||
|
|
import torch.nn.functional as F
|
|||
|
|
|
|||
|
|
__all__ = ["moe_forward", "mixtral_decoder_layer_forward"]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _flatten(x):
|
|||
|
|
"""[B,S,H] -> [B*S,H]，返回还原所需的形状；已是 2D 则原样返回。"""
|
|||
|
|
if x.dim() == 3:
|
|||
|
|
return x.reshape(-1, x.shape[-1]), x.shape[:2]
|
|||
|
|
return x, None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def moe_forward(hidden_states, router_logits, experts, top_k=2,
|
|||
|
|
renormalize=True, **kwargs):
|
|||
|
|
"""通用 MoE forward。
|
|||
|
|
|
|||
|
|
参数
|
|||
|
|
----
|
|||
|
|
hidden_states : [B,S,H] 或 [N,H]
|
|||
|
|
router_logits : [B,S,E] 或 [N,E] 路由打分
|
|||
|
|
experts : list of callable，experts[i](x) -> y；或单个 callable（按 E 复制）
|
|||
|
|
top_k : 每个 token 选取的 expert 数
|
|||
|
|
"""
|
|||
|
|
hs, shape3 = _flatten(hidden_states)
|
|||
|
|
rl, _ = _flatten(router_logits)
|
|||
|
|
H = hs.shape[-1]
|
|||
|
|
|
|||
|
|
if not isinstance(experts, (list, tuple)):
|
|||
|
|
experts = [experts] * rl.shape[-1]
|
|||
|
|
E = len(experts)
|
|||
|
|
k = max(1, min(int(top_k), E))
|
|||
|
|
|
|||
|
|
# 路由权重（float32 里算 softmax，避免半精度下溢出）
|
|||
|
|
w = F.softmax(rl.float(), dim=-1)
|
|||
|
|
top_w, top_i = torch.topk(w, k, dim=-1)
|
|||
|
|
if renormalize:
|
|||
|
|
top_w = top_w / top_w.sum(dim=-1, keepdim=True).clamp_min(1e-20)
|
|||
|
|
top_w = top_w.to(hs.dtype)
|
|||
|
|
|
|||
|
|
out = torch.zeros_like(hs)
|
|||
|
|
for j in range(k):
|
|||
|
|
idx = top_i[:, j]
|
|||
|
|
wj = top_w[:, j].unsqueeze(-1)
|
|||
|
|
for e in range(E):
|
|||
|
|
mask = idx == e
|
|||
|
|
if not bool(mask.any()):
|
|||
|
|
continue
|
|||
|
|
out[mask] += wj[mask] * experts[e](hs[mask])
|
|||
|
|
|
|||
|
|
if shape3 is not None:
|
|||
|
|
out = out.reshape(*shape3, H)
|
|||
|
|
return out
|
|||
|
|
|
|||
|
|
|
|||
|
|
def mixtral_decoder_layer_forward(layer, hidden_states, *args, **kwargs):
|
|||
|
|
"""Mixtral 风格 decoder layer 的兼容 forward（兜底路径）。
|
|||
|
|
|
|||
|
|
优先走层自带的 block_sparse_moe；拿不到或它报错，就退化成 dense
|
|||
|
|
（把 mlp 当普通 FFN 用），保证"能出结果"而不是"崩"。
|
|||
|
|
"""
|
|||
|
|
residual = hidden_states
|
|||
|
|
h = layer.input_layernorm(hidden_states)
|
|||
|
|
h = layer.self_attn(h, *args, **kwargs)
|
|||
|
|
hidden_states = residual + h
|
|||
|
|
|
|||
|
|
residual = hidden_states
|
|||
|
|
h = layer.post_attention_layernorm(hidden_states)
|
|||
|
|
moe = getattr(layer, "block_sparse_moe", None) or getattr(layer, "mlp", None)
|
|||
|
|
if moe is None:
|
|||
|
|
raise AttributeError("layer 上既没有 block_sparse_moe 也没有 mlp")
|
|||
|
|
try:
|
|||
|
|
return residual + moe(h)
|
|||
|
|
except Exception: # noqa: BLE001 原生 MoE 路径失败 -> 退化 dense
|
|||
|
|
out = moe(h)
|
|||
|
|
if isinstance(out, tuple):
|
|||
|
|
out = out[0]
|
|||
|
|
return residual + out
|