基于 Amazon EKS、EFA 与 DeepEP 的 MoE 强化学习大规模扩展与吞吐量优化

[ 跨节点 MoE Token 路由示意图 ]

+-----------------------+                    +-----------------------+
| 节点 1 (GPU 0..7)     |                    | 节点 2 (GPU 0..7)     |
| - Token 生成节点       |                    | - 专家 5..8           |
| - 专家 1..4           |                    |                       |
+-----------+-----------+                    +-----------+-----------+
            |                                            ^
            |-------- 跨节点 All-to-All 通信 -------------|
                      (AWS EFA + DeepEP 内核)

# eksctl 节点组配置文件示例
apiVersion: eksctl.io/v1alpha5
kind: ClusterConfig
metadata:
  name: moe-rl-eks-cluster
  region: us-east-1

managedNodeGroups:
  - name: p5-efa-nodes
    instanceType: p5.48xlarge
    minSize: 4
    maxSize: 16
    desiredCapacity: 8
    volumeSize: 500
    privateNetworking: true
    efa:
      enabled: true
    iam:
      withAddonPolicies:
        autoScaler: true

apiVersion: apps/v1
kind: DaemonSet
metadata:
  name: aws-efa-k8s-device-plugin
  namespace: kube-system
spec:
  selector:
    matchLabels:
      name: aws-efa-k8s-device-plugin
  template:
    metadata:
      labels:
        name: aws-efa-k8s-device-plugin
    spec:
      hostNetwork: true
      containers:
        - image: 602401143452.dkr.ecr.us-east-1.amazonaws.com/aws-efa-k8s-device-plugin:v0.5.5
          name: aws-efa-k8s-device-plugin
          securityContext:
            allowPrivilegeEscalation: false
            capabilities:
              drop: ["ALL"]
          volumeMounts:
            - name: device-plugin
              mountPath: /var/lib/kubelet/device-plugins
      volumes:
        - name: device-plugin
          hostPath:
            path: /var/lib/kubelet/device-plugins

import torch
import torch.distributed as dist
# 引入编译了 EFA/RDMA 支持的 deep_ep 模块
import deep_ep

class OptimizedMoERouter(torch.nn.Module):
    def __init__(self, num_experts: int, hidden_dim: int, top_k: int):
        super().__init__()
        self.num_experts = num_experts
        self.hidden_dim = hidden_dim
        self.top_k = top_k
        
        # 初始化 DeepEP 缓冲区，自动管理 NVLink 和 EFA RDMA 传输
        self.ep_buffer = deep_ep.Buffer(
            group=dist.group.WORLD,
            num_nvl_bytes=1024 * 1024 * 128,  # 节点内 NVLink 缓冲区
            num_rdma_bytes=1024 * 1024 * 512   # 节点间 EFA RDMA 缓冲区
        )

    def forward(self, hidden_states: torch.Tensor, router_logits: torch.Tensor):
        # 计算路由权重与目标专家 ID
        weights, selected_experts = torch.topk(
            torch.softmax(router_logits, dim=-1), 
            self.top_k, 
            dim=-1
        )
        
        # 使用 DeepEP 内核进行 Token 分发（Dispatch）
        # 避开了传统 NCCL 的高延迟开销
        dispatched_tokens, handle = self.ep_buffer.dispatch(
            hidden_states, 
            selected_experts
        )
        
        # 在分配的 GPU 上执行专家计算
        expert_outputs = self.compute_experts(dispatched_tokens)
        
        # 聚合阶段（Combine）：将计算结果按原路径传回
        combined_output = self.ep_buffer.combine(expert_outputs, handle)
        
        return combined_output

    def compute_experts(self, tokens: torch.Tensor) -> torch.Tensor:
        # 专家网络计算逻辑占位符
        return tokens * 1.05