# 专家并行 (EP)

专家并行将混合专家模型（MoE）的专家分布到多个 GPU 上，每个 rank 只持有部分专家。这降低了单卡显存占用，使大规模 MoE 模型的训练成为可能。

## 概览

| 概念 | 说明 |
|------|------|
| **ExpertParallelConfig** | 控制 EP 行为的配置数据类 |
| **apply_expert_parallel()** | 入口函数，负责分片专家并替换前向传播 |
| **shard_experts()** | 将专家均匀分配到各 EP rank |
| **patch_forward()** | 将 MoE block 的 forward 替换为带 all-to-all 通信的 EP 版本 |

## 配置

```python
from twinkle.model.transformers.moe.expert_parallel import ExpertParallelConfig

config = ExpertParallelConfig(
    enabled=True,              # 启用专家并行
    router_dtype='fp32',       # 路由计算精度：'fp32', 'bf16', 'fp16'
    keep_router_logits=True,   # 在输出中保留路由 logits
    ignore_shared_experts=False,# 跳过共享专家计算（如 DeepSeek）
    ep_size=None,              # EP 并行度（由 TransformersModel 使用）
)
```

## 配合 DeviceMesh 使用

在 `DeviceMesh.from_sizes()` 中设置 `ep_size` 即可激活 EP。框架会在模型初始化时自动调用 `apply_expert_parallel()`。

```python
from twinkle.utils import DeviceMesh

# 8 卡：2 路 EP × 4 路数据并行
device_mesh = DeviceMesh.from_sizes(
    world_size=8,
    dp_size=4,
    ep_size=2,
)
```

EP + FSDP 组合分片：

```python
# 8 卡：2 路 EP，每个 EP 组内 2 路 FSDP
device_mesh = DeviceMesh.from_sizes(
    world_size=8,
    dp_size=2,
    ep_size=2,
    ep_fsdp_size=2,
)
```

## 通信模式

EP 前向传播遵循 4 阶段流水线：

1. **预处理** — 计算每个专家的 token 数量和分割大小
2. **Token Pre-All2All** — 按专家分配排列 token，然后在 EP rank 间执行 all-to-all 交换
3. **专家计算** — 每个 rank 在接收到的 token 上运行本地专家
4. **Token Post-All2All** — all-to-all 交换结果，反排列并应用路由权重

```
输入 token → 路由器 → [预处理] → [pre_all2all] → [本地专家] → [post_all2all] → 输出
```

## 要求

- `num_experts` 必须能被 `ep_size` 整除
- `torch.distributed` 必须已初始化
- MoE block 必须定义 `gate`/`router` 模块和 `experts`（支持 `nn.ModuleList` 或张量形式的 `gate_up_proj`/`down_proj`）
- 共享专家（如 DeepSeek MoE）会自动处理，除非设置 `ignore_shared_experts=True`