einsum函数

理解专家并行，需要了解einsum函数

import torch# 设置输入张量的维度：s = 3 tokens, e = 2 experts, c = 2 capacity, m = 4 embedding dim
s, e, c, m = 3, 2, 2, 4# 1. 输入 token 的嵌入向量 (s, m)
reshaped_input = torch.tensor([[1.0, 1.0, 1.0, 1.0],  # token 0[2.0, 2.0, 2.0, 2.0],  # token 1[3.0, 3.0, 3.0, 3.0],  # token 2
])# 2. dispatch_mask: (s, e, c)
# 表示每个 token 被分配到哪个 expert 的哪个槽位（slot）
dispatch_mask = torch.tensor([# token 0[[1, 0],   # expert 0: slot 0[0, 0]],  # expert 1: no slot# token 1[[0, 0],[1, 0]],  # expert 1: slot 0# token 2[[0, 1],   # expert 0: slot 1[0, 0]],  # expert 1: no slot
])
dispatch_mask = dispatch_mask.float()
# 3. 应用 einsum 进行 token 分发到专家
dispatched_input = torch.einsum("sec,sm->ecm", dispatch_mask, reshaped_input)# 4. 打印结果
print("Dispatched Input shape:", dispatched_input.shape)
print("\nDispatched Input Tensor:")
print(dispatched_input)

结果如下

Dispatched Input shape: torch.Size([2, 2, 4])Dispatched Input Tensor:
tensor([[[1., 1., 1., 1.],[3., 3., 3., 3.]],[[2., 2., 2., 2.],[0., 0., 0., 0.]]])

增加画图功能

import torch
import matplotlib.pyplot as plt
# 输入数据
reshaped_input = torch.tensor([[1.0, 1.0, 1.0, 1.0],[2.0, 2.0, 2.0, 2.0],[3.0, 3.0, 3.0, 3.0],
])  # float32dispatch_mask = torch.tensor([[[1, 0], [0, 0]],[[0, 0], [1, 0]],[[0, 1], [0, 0]],
])  # int64 → 不兼容# 修复：转换为 float 类型
dispatch_mask = dispatch_mask.float()# Einsum 分发
dispatched_input = torch.einsum("sec,sm->ecm", dispatch_mask, reshaped_input)# 输出
print("Dispatched input shape:", dispatched_input.shape)
print(dispatched_input)def visualize_dispatch(dispatch_mask):s, e, c = dispatch_mask.shape  # tokens, experts, capacityplt.figure(figsize=(6, 4))for token in range(s):for expert in range(e):for slot in range(c):if dispatch_mask[token, expert, slot] > 0:# token 位置 (左边)x_token, y_token = 0, s - token# expert-slot 位置 (右边)x_expert, y_expert = 4, e * c - (expert * c + slot)# 画连接线plt.plot([x_token, x_expert], [y_token, y_expert], 'k-', lw=1)# 标记 tokenplt.text(x_token - 0.2, y_token, f"T{token}", va='center', ha='right', fontsize=10)# 标记 expert-slotplt.text(x_expert + 0.2, y_expert, f"E{expert}-S{slot}", va='center', ha='left', fontsize=10)# 设置图形样式plt.xlim(-1, 6)plt.ylim(0, max(s, e*c) + 1)plt.axis('off')plt.title("Token → Expert-Slot Routing")plt.show()visualize_dispatch(dispatch_mask)