应用代码
python
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
# ============================================================
# 模块 1：字节级分词器
# ============================================================
class ByteTokenizer:
def __init__(self):
self.vocab_size = 259
self.bos = 256 # 序列开始，告诉LLM一个独立的文本片段或输入样本从这里开始。
self.eos = 257 # 序列结束，告诉LLM一个文本片段到这里结束。
# <pad>会被添加到较短序列的末尾使批次中所有序列长度一致，便于高效的矩阵运算。
self.pad = 258 # 填充，在模型训练或推理时，通常需要将多条长短不一的文本组成一个批次。
'''
1. 基础字节编码：数量256个，它们代表了计算机中所有可能的单字节值从0到255。这种方法能够确保任何文本，不论其语言或编码，都能被无损地编码成一串数字Token ID
2. 特殊功能编码：数量3个，这些Token专门用于提供文本结构信息，确保模型能够正确处理和理解文本段落的边界和批处理时的对齐，便于计算。
'''
def encode(self, text, add_bos=True, add_eos=True):
# 把输入文本信息进行utf-8字符集编码，得到字节序列b
# 每个字节值0-255对应一个Token ID
b = text.encode("utf-8", errors="surrogatepass")
ids = list(b) # 将UTF-8字节序列转换为Token ID列表
if add_bos:
# 标记文本开头，添加<bos> Token ID
ids = [self.bos] + ids
if add_eos:
# 标记文本结束，添加<eos> Token ID
ids = ids + [self.eos]
return ids
def decode(self, ids):
# 过滤特殊 token
ids = [i for i in ids if i < 256]
return bytes(ids).decode("utf-8", errors="replace")
def batch_encode(self, texts, pad_to=None):
# pad_to 用于规定批处理中每条 Token ID 序列的目标长度（强制对齐）
encs = [self.encode(t) for t in texts]
# 如果未指定 pad_to，则使用当前批次中最长序列的长度；否则使用 pad_to 规定的长度
maxlen = max(len(x) for x in encs) if pad_to is None else pad_to
# 将所有序列填充到maxlen长度，填充方式是在每条序列末尾添加[pad]，强制对齐形成规则的张量
# batch_encode阶段返回对齐处理张量、未对齐处理的序列长度的考虑：
# 不规则的张量不能直接输入到为高性能并行计算优化的硬件GPU、TPU中
# 对齐是进行批处理和利用硬件并行性的必要预处理步骤。这种填充虽然解决了并行计算的问题
# 但也引入了计算冗余比如这里的[pad]。
# 原始序列长度信息，则是为了告诉模型末尾填充[pad]从哪里开始的
# 从而在Attention机制中屏蔽掉它们，防止其将计算资源和注意力分散到这些无关紧要的数据上
# 确保模型只关注真实的输入信息。
arr = [x + [self.pad] * (maxlen - len(x)) for x in encs]
# 记录原始序列的真实长度，这条信息将用于Attention，避免模型关注到[pad] Token
lengths = torch.LongTensor([len(x) for x in encs])
# 经过填充对齐Token ID张量输入给模型，原始序列的真实长度张量用于Attention
return torch.LongTensor(arr), lengths
# ============================================================
# 模块 2：RoPE 旋转位置编码
# ============================================================
class RoPE(nn.Module):
def __init__(self, d_k, max_len=4096, base=10000.0):
super().__init__()
# 频率：theta_i = base^(-2i/d_k)
inv_freq = 1.0 / (base ** (torch.arange(0, d_k, 2).float() / d_k))
self.register_buffer("inv_freq", inv_freq)
self.max_len = max_len
def forward(self, x, offset=0):
"""
x: [B, H, T, d_k]
offset: KV Cache 中已有长度，用于生成正确的绝对位置
"""
B, H, T, d_k = x.shape
positions = torch.arange(offset, offset + T, device=x.device).float()
freqs = torch.outer(positions, self.inv_freq) # [T, d_k/2]
emb = torch.cat([freqs, freqs], dim=-1) # [T, d_k]
cos = emb.cos()[None, None, :, :] # [1,1,T,d_k]
sin = emb.sin()[None, None, :, :]
# 把 x 的相邻两维看成一对 (x1, x2)，做二维旋转
x1 = x[..., 0::2]
x2 = x[..., 1::2]
rotated = torch.stack(
[x1 * cos[..., 0::2] - x2 * sin[..., 0::2],
x1 * sin[..., 1::2] + x2 * cos[..., 1::2]],
dim=-1
).flatten(-2)
return rotated
# ============================================================
# 模块 3：带 RoPE 和 KV Cache 的自注意力
# ============================================================
class SelfAttentionWithRoPEKV(nn.Module):
def __init__(self, d_model, nhead):
super().__init__()
assert d_model % nhead == 0
self.nhead = nhead
self.d_k = d_model // nhead
self.qkv = nn.Linear(d_model, d_model * 3)
self.out = nn.Linear(d_model, d_model)
self.rope = RoPE(self.d_k)
def forward(self, x, mask=None, kv_cache=None, use_cache=False):
"""
x: [B, T, D]
kv_cache: dict 或 None，包含 "k" 和 "v"，形状 [B, H, T_past, d_k]
use_cache: 是否返回新的 cache
"""
B, T, D = x.shape
qkv = self.qkv(x)
q, k, v = qkv.chunk(3, dim=-1)
q = q.view(B, T, self.nhead, self.d_k).transpose(1, 2)
k = k.view(B, T, self.nhead, self.d_k).transpose(1, 2)
v = v.view(B, T, self.nhead, self.d_k).transpose(1, 2)
# KV Cache：拼接历史 K、V
past_len = 0
if kv_cache is not None:
past_len = kv_cache["k"].shape[2]
k = torch.cat([kv_cache["k"], k], dim=2)
v = torch.cat([kv_cache["v"], v], dim=2)
# 对 Q、K 应用 RoPE，位置从 past_len 开始
q = self.rope(q, offset=past_len)
k = self.rope(k, offset=past_len)
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
# 因果掩码：query i 只能看 key j <= i
total_len = past_len + T
causal = torch.tril(torch.ones(T, total_len, device=x.device)).bool()
scores = scores.masked_fill(~causal[None, None, :, :], float("-inf"))
# padding 掩码
if mask is not None:
# mask: [B, total_len]，只对当前新 token 部分做 padding 屏蔽
pad_mask = mask[:, None, None, :].bool()
scores = scores.masked_fill(~pad_mask, float("-inf"))
attn = F.softmax(scores, dim=-1)
out = torch.matmul(attn, v)
out = out.transpose(1, 2).contiguous().view(B, T, D)
out = self.out(out)
new_cache = None
if use_cache:
new_cache = {"k": k, "v": v}
return out, new_cache
# ============================================================
# 模块 4：MoE 层（原文 MoELayer，保留噪声门控与容量限制）
# ============================================================
class MoELayer(nn.Module):
def __init__(self, d_model, d_ff, n_experts=4, k=1,
capacity_factor=1.25, noisy_gating=True):
super().__init__()
assert k in (1, 2) # 确保K激活专家数是1或2
self.n_experts = n_experts
self.k = k
self.capacity_factor = capacity_factor
self.noisy_gating = noisy_gating
# 门控网络：负责计算每个token与n_experts个专家的匹配得分（Logits）。
self.w_gating = nn.Linear(d_model, n_experts, bias=False)
if noisy_gating:
# 噪声网络：引入噪声有助于在训练时平均分配token到不同的专家，缓解负载不均衡问题。
self.w_noise = nn.Linear(d_model, n_experts, bias=False)
# 专家网络，每个专家是一个独立的FFN
self.experts = nn.ModuleList([
nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(), # 使用GELU激活函数
nn.Linear(d_ff, d_model)
) for _ in range(n_experts)
])
def _noisy_logits(self, x):
"""
x : 展平后的输入token向量，形状为[N, D] (N=B*T)。
Tensor: 带噪声的专家logits，形状为[N, E]。
"""
logits = self.w_gating(x)
# 在训练模式且开启noisy_gating时引入随机噪声
if self.noisy_gating and self.training:
# 使用sigmoid将w_noise的输出映射到[0, 1]，作为噪声的标准差
noise_std = torch.sigmoid(self.w_noise(x))
# 加上正态分布噪声这增强了随机性，有助于训练时的负载均衡
logits = logits + torch.randn_like(logits) * noise_std
return logits
def forward(self, x):
B, T, D = x.shape
N = B * T
x_flat = x.view(N, D)
logits = self._noisy_logits(x_flat)
scores = F.softmax(logits, dim=-1) # 归一化的专家选择权重，[N, E]
if self.k == 1:
top1 = torch.argmax(scores, dim=-1) # 每个token选出的Top-1专家索引，[N]
# Dispatch Mask：[N, E]，标记每个token选中的Top-1专家，用1表示选中
dispatch = F.one_hot(top1, num_classes=self.n_experts).to(x.dtype)
# 提取每个token选中的Top-1专家的得分作为最终组合权重，得到[N]
combine_w = torch.gather(scores, 1, top1.unsqueeze(1)).squeeze(1)
# 计算规定每个专家最多处理的token数量
capacity = int((N / self.n_experts) * self.capacity_factor) + 1
# 初始化输出
out_flat = torch.zeros_like(x_flat)
for e in range(self.n_experts):
# 找到专家e应该处理的token原始索引值，[N]
idx = torch.nonzero(dispatch[:, e], as_tuple=False).squeeze(-1)
if idx.numel() > capacity:
# 专家e容量检查，如果超过容量，丢弃多余的token
idx = idx[:capacity]
if idx.numel() == 0:
continue
# 第e个专家处理token
y = self.experts[e](x_flat[idx])
out_flat[idx] = y
# 所有专家处理的结果乘以组合权重
out_flat = out_flat * combine_w.unsqueeze(1)
return out_flat.view(B, T, D)
else:
# Top-2 简化实现
# 每个token选出Top-2专家的得分和索引，[N, 2]
topk_vals, topk_idx = torch.topk(scores, k=2, dim=-1)
# 计算出每个专家最大处理token数量
capacity = int((N / self.n_experts) * self.capacity_factor) + 1
# 初始化存储空间
buckets = [[] for _ in range(self.n_experts)]
for i in range(N):
for j in range(2):
e = int(topk_idx[i, j].item()) # Top-K专家索引值
w = float(topk_vals[i, j].item()) # 相应的token组合权重
buckets[e].append((i, w)) # 存储：token原始索引、权重
out_flat = torch.zeros_like(x_flat) # 初始化输出结果
for e in range(self.n_experts):
bucket = buckets[e]
if not bucket:
continue
bucket = bucket[:capacity] # 每个专家丢弃超过容量的token
# token的原始索引值转化为张量：[C] ->（C=容量限制后的数量）
idxs = torch.tensor([i for i, _ in bucket], device=x.device)
# 对应的组合权重转化为张量：[C]
weights = torch.tensor([w for _, w in bucket], device=x.device)
y = self.experts[e](x_flat[idxs]) # 获取专家e需要处理的token，[C, D]
# 专家的输出乘以权重后，累加到输出张量上（Top-2叠加），可能有不止一个专家处理同一个token
out_flat[idxs] += y * weights.unsqueeze(1)
return out_flat.view(B, T, D)
# ============================================================
# 模块 5：Transformer Block
# ============================================================
class TransformerBlock(nn.Module):
def __init__(self, d_model, nhead, d_ff, use_moe=False,
moe_params=None, dropout=0.1):
super().__init__()
# 第一子层：多头自注意力机制
self.attn = SelfAttentionWithRoPEKV(d_model, nhead)
# Layer Normalization层：LN1位于注意力层之前
self.ln1 = nn.LayerNorm(d_model)
# Layer Normalization层：LN2位于FFN、MoE层之前
self.ln2 = nn.LayerNorm(d_model)
# Dropout层
self.dropout = nn.Dropout(dropout)
self.use_moe = use_moe
# 第二子层：根据use_moe决定使用FFN还是MoE
if use_moe:
assert moe_params is not None
# 稀疏MoE层
self.moe = MoELayer(**moe_params)
else:
# 传统的前馈网络(FFN)
self.ffn = nn.Sequential(
nn.Linear(d_model, d_ff), # 扩展维度
nn.GELU(), # 激活函数
nn.Linear(d_ff, d_model) # 还原维度
)
def forward(self, x, mask=None, kv_cache=None, use_cache=False):
# 第一子层：LayerNorm → Attention(RoPE + KV Cache) → 残差
attn_out, new_cache = self.attn(
self.ln1(x), mask=mask, kv_cache=kv_cache, use_cache=use_cache
)
x = x + self.dropout(attn_out)
# 第二子层：LayerNorm → FFN/MoE → 残差
if self.use_moe:
moe_out = self.moe(self.ln2(x))
x = x + self.dropout(moe_out)
else:
ffn_out = self.ffn(self.ln2(x))
x = x + self.dropout(ffn_out)
return x, new_cache
# ============================================================
# 模块 6：完整 MoE + LLM
# ============================================================
class MiniMoELLM(nn.Module):
def __init__(self, vocab_size, d_model=64, nhead=4, n_layers=2,
d_ff=256, use_moe_layer_index=None, moe_params=None):
"""
use_moe_layer_index: 哪些层使用MoE，例如[1,3]
moe_params: MoE参数字典，会自动注入 d_model和d_ff
"""
super().__init__()
self.vocab_size = vocab_size # 词汇表大小，不用考虑特殊token的预测
self.d_model = d_model # Token Embedding 维度
self.tok_emb = nn.Embedding(vocab_size, d_model) # Token嵌入层
# 注意：RoPE 已经在注意力内部编码位置，这里不需要额外位置编码
# 判断是否使用MoE
if use_moe_layer_index is None:
use_moe_layer_index = set() # 默认使用标准FFN
else:
use_moe_layer_index = set(use_moe_layer_index)
# 配置MoE相关参数
if moe_params is not None:
moe_params = moe_params.copy() # 复制参数，注入LLM的d_model和d_ff
moe_params.setdefault("d_model", d_model)
moe_params.setdefault("d_ff", d_ff)
# Transformer层
self.layers = nn.ModuleList()
for i in range(n_layers):
use_moe = i in use_moe_layer_index # 确定当前层是否使用MoE模块
self.layers.append(
TransformerBlock(
d_model=d_model, nhead=nhead, d_ff=d_ff,
use_moe=use_moe, moe_params=moe_params
)
)
# LayNorm+输出层，共享embedding权重
self.ln_f = nn.LayerNorm(d_model) # 最终Layer Normalization
self.lm_head = nn.Linear(d_model, vocab_size, bias=False) # 语言模型头，logits投影
self.lm_head.weight = self.tok_emb.weight # 权重共享
def forward(self, idx, mask=None, kv_caches=None, use_cache=False):
"""
idx: [B, T]
kv_caches: list of dict，长度等于层数；None 表示没有缓存
"""
B, T = idx.shape
x = self.tok_emb(idx) # 输入嵌入=Token Embedding
new_caches = [] if use_cache else None
for i, blk in enumerate(self.layers):
cache_i = kv_caches[i] if kv_caches is not None else None
# 经过Transformer块，包含Attention和FFN、MoE
x, new_cache_i = blk(
x, mask=mask, kv_cache=cache_i, use_cache=use_cache
)
if use_cache:
new_caches.append(new_cache_i)
x = self.ln_f(x) # 最终的层归一化
logits = self.lm_head(x) # 投影到词汇表维度，得到logits[B, T, vocab_size]
# 返回Logits，用于损失计算或Softmax后的概率预测
return logits, new_caches
# ============================================================
# 模块 7：训练与生成示例
# ============================================================
def train_and_generate():
torch.manual_seed(42)
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = ByteTokenizer()
# 极小语料
corpus = [
"MoE is powerful.",
"Mixture of experts is efficient.",
"Deep learning is fun.",
"语言模型很强大。",
"专家混合模型很高效。",
]
# 训练数据：每条文本 -> input/target 错一位
input_ids, lengths = tokenizer.batch_encode(corpus)
# .to(device) 是 PyTorch 中把张量（Tensor）或模型（Module）移动到指定设备的方法。
# 这里的 device 通常是 "cuda"（GPU）或 "cpu"。
input_ids = input_ids.to(device)
targets = input_ids.clone()
targets[:, :-1] = input_ids[:, 1:]
targets[:, -1] = tokenizer.pad
# 构建模型：第 1 层使用 MoE，第 0 层使用普通 FFN
model = MiniMoELLM(
vocab_size=tokenizer.vocab_size,
d_model=64,
nhead=4,
n_layers=2,
d_ff=256,
use_moe_layer_index=[1],
moe_params={
"n_experts": 4,
"k": 1,
"capacity_factor": 1.25,
"noisy_gating": True,
},
).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
# 训练
model.train()
for step in range(200):
optimizer.zero_grad()
logits, _ = model(input_ids)
loss = F.cross_entropy(
logits.view(-1, tokenizer.vocab_size),
targets.view(-1),
ignore_index=tokenizer.pad
)
loss.backward()
optimizer.step()
if step % 50 == 0:
print(f"step {step:3d} loss {loss.item():.4f}")
# ============================================================
# 生成：使用 KV Cache
# ============================================================
model.eval()
prompt = "MoE"
ids = tokenizer.encode(prompt, add_bos=True, add_eos=False)
input_tensor = torch.LongTensor([ids]).to(device)
generated = ids.copy()
kv_caches = None
with torch.no_grad():
# 第一步：处理整个 prompt，建立 KV Cache
logits, kv_caches = model(
input_tensor, kv_caches=None, use_cache=True
)
next_id = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True) # [1,1]
generated.append(next_id.item())
# 后续步骤：每次只输入新 token，复用 KV Cache
for _ in range(40):
logits, kv_caches = model(
next_id, kv_caches=kv_caches, use_cache=True
)
next_id = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
generated.append(next_id.item())
if next_id.item() == tokenizer.eos:
break
print("\n生成结果：")
print(tokenizer.decode(generated))
if __name__ == "__main__":
train_and_generate()整体结构
text
ByteTokenizer
→ Token Embedding
→ N × TransformerBlock
├─ LayerNorm → SelfAttention(RoPE + KV Cache) → 残差
└─ LayerNorm → MoELayer（Top-K 专家）→ 残差
→ LayerNorm
→ lm_head（权重共享）
→ logits生成时，SelfAttention 使用 KV Cache，每步只计算新 token 的 Q、K、V。
关键模块说明
RoPE 的作用
python
# ============================================================
# 模块 2：RoPE 旋转位置编码
# ============================================================
class RoPE(nn.Module):
def __init__(self, d_k, max_len=4096, base=10000.0):
super().__init__()
# 频率：theta_i = base^(-2i/d_k)
inv_freq = 1.0 / (base ** (torch.arange(0, d_k, 2).float() / d_k))
self.register_buffer("inv_freq", inv_freq)
self.max_len = max_len
def forward(self, x, offset=0):
"""
x: [B, H, T, d_k]
offset: KV Cache 中已有长度，用于生成正确的绝对位置
"""
B, H, T, d_k = x.shape
positions = torch.arange(offset, offset + T, device=x.device).float()
freqs = torch.outer(positions, self.inv_freq) # [T, d_k/2]
emb = torch.cat([freqs, freqs], dim=-1) # [T, d_k]
cos = emb.cos()[None, None, :, :] # [1,1,T,d_k]
sin = emb.sin()[None, None, :, :]
# 把 x 的相邻两维看成一对 (x1, x2)，做二维旋转
x1 = x[..., 0::2]
x2 = x[..., 1::2]
rotated = torch.stack(
[x1 * cos[..., 0::2] - x2 * sin[..., 0::2],
x1 * sin[..., 1::2] + x2 * cos[..., 1::2]],
dim=-1
).flatten(-2)
return rotated- 对 Q、K 的相邻维度对做旋转；
- 旋转角度与绝对位置有关；
- offset 参数在 KV Cache 生成时非常关键：新 token 的位置从 past_len 开始，保证位置编码连续。
KV Cache 的作用
python
if kv_cache is not None:
past_len = kv_cache["k"].shape[2]
k = torch.cat([kv_cache["k"], k], dim=2)
v = torch.cat([kv_cache["v"], v], dim=2)- 生成第 ttt 个 token 时，只计算新 token 的 Q、K、V；
- 历史 K、V 从缓存读取，拼接后参与注意力；
- 注意力分数矩阵形状从 B,H,1,1 变成 B,H,1,past_len+1，只算一行；
- 计算量从 O(L2)O(L^2)O(L2) 降到 O(L)O(L)O(L) 每步。
因果掩码与 RoPE 配合
python
causal = torch.tril(torch.ones(T, total_len)).bool()
scores = scores.masked_fill(~causal[None,None,:,:], float("-inf"))- 有 KV Cache 时，total_len = past_len + T；
- 新 token 的 query 只能看历史 key 和当前 key；
- RoPE 的 offset=past_len 保证位置编码与绝对位置一致。
MoE 层
保留原文的：
- w_gating 路由打分；
- w_noise 噪声门控；
- Top-1/Top-2 选择；
- capacity 容量限制；
- 专家 FFN 加权合并。
总结
这个完整例子包含：
- 字节级分词器：任意文本 → 字节 ID；
- RoPE：对 Q、K 做位置相关旋转，支持相对位置；
- KV Cache：生成时只算新 token 的 K、V，历史复用；
- 自注意力：RoPE + 因果掩码 + KV Cache；
- MoE 层：Top-K 路由 + 噪声门控 + 容量限制；
- Transformer Block：注意力 + MoE + 残差 + LayerNorm；
- MiniMoELLM：Embedding + N 层 Block + ln_f + lm_head；
- 训练与生成：交叉熵训练，带 KV Cache 的自回归生成。
代码可直接运行，训练后能续写训练语料中的简单句子，完整展示了现代 MoE-LLM 的核心组件如何协同工作。