MegaScale

• 算法优化
  • 并行 transformer：原版 transformer 公式展开可以让 attention 块和 MLP 块并行计算，提高速度
  • 滑动窗口 attention
  • LAMB 优化器
• 3D 并行中的通信重叠
  • 数据并行中的重叠：初始的 all-gather 操作会在每一轮迭代开始的时候随着数据读入进行预取，完成重叠
  • 流水线并行中的重叠：把 send/recieve 操作解耦，send 和 recieve 可以异步执行来重叠
• 提效操作
  • 数据流水线
    • 异步的数据预处理：数据预处理不在关键路径上，因此可以在每一轮同步梯度的时候，为下一步准备的数据预处理就可以开始了
    • 传统的分布式训练中，每个 worker 有自己的 data loader，这导致了 workers 之间会竞争硬盘的读取带宽，从而造成瓶颈。然而同一个机器上的多个 GPU 通常属于同一个张量并行组，他们的数据应该是一样的，因此我们在每台机器上使用单个专用数据加载器将训练数据读取到一块共享内存中。
  • 通信组优化
    • 把 pytorch 自带的 KV 存储换成了 redis，减少了初始化的耗时
    • 优化了通信组的初始化顺序，最小化了全局屏障的影响
• 容灾
  • 用 k8s 来管理，并且设计了一套自查的系统，当碰到异常或者心跳丢失，会通知所有的执行程序进行自查，以定位异常的节点。然后将所有异常的节点剔除，加入等量的健康节点，然后重新读取 checkpoint，重启训练
  • 同时 RDAM 监控会用作监控网络利用率和效率，以发现那些隐秘的异常
  • 快速的 checkpoint 和恢复
    • 第一阶段：片上状态存到 host 的内存里
    • 第二阶段：后台进程一步下刷到分布式文件系统中（HDFS）
    • 读取的时候，观察到一个数据并行组的所有 GPU 会从同一个 HDFS 中读取相同的模型状态，存在带宽瓶颈，这可以优化成由其中一个读取，然后广播给其他的