AI并行计算的奥秘：揭秘DP、PP、TP、EP

AI计算，尤其是模型训练和推理，主要依赖并行计算。

AI算法如矩阵相乘、卷积等，需要借助成千上万的GPU，通过并行任务来缩短计算时间。

搭建并行计算框架，通常使用以下并行方式：

• Data Parallelism（数据并行）
• Pipeline Parallelism（流水线并行）
• Tensor Parallelism（张量并行）
• Expert Parallelism（专家并行）

接下来，我们逐一探讨这些并行计算方式的工作原理。

数据并行（DP）

数据并行是AI训练中常见的并行方式，核心思想是每个GPU都拥有完整的模型副本，训练数据被分割成多个小批次，每个批次分配给不同的GPU。

数据并行的过程包括：

1. 均匀切割数据，发给不同GPU；
2. 各GPU独立进行前向传播、反向传播，计算梯度；
3. 通过All-Reduce通信方式，将梯度汇总到Server GPU；
4. Server GPU将全局梯度广播回每个Worker GPU，更新参数。

数据并行的优点在于实现简单，能显著加速大规模数据训练。缺点是显存限制和通信开销大。

ZeRO优化器

ZeRO（Zero Redundancy Optimizer）通过切分模型副本，减少内存占用。ZeRO有3个阶段：

• ZeRO-1：优化器状态划分；
• ZeRO-2：优化器状态和梯度划分；
• ZeRO-3：优化器状态、梯度和参数全面划分。

流水线并行（PP）

流水线并行将模型的不同层分配到不同GPU上，按顺序处理数据。虽然看似串行，但通过micro-batch和提前调度，能减少空闲时间。

张量并行（TP）

张量并行将模型的张量按维度切分到不同GPU上。优点是适合大张量，减少内存占用；缺点是通信开销大。

专家并行

专家并行是MoE（混合专家模型）的一种并行策略，通过动态路由机制分配专家处理任务。需要合理的门控机制和专家选择策略。

混合并行

实际应用中，常结合多种并行策略。例如数据并行+张量并行、流水线并行+专家并行等。

总结

了解这些并行计算方式有助于理解算力集群架构和网络设计。不同方式有不同的通信流量特点，设计需适配这些特点以提高效率。