torchgpipe — torchgpipe 0.0.7 documentation

目录

1. GPipe Pipeline 概述
2. 核心模块解析
   • Fork 和 Join 的作用
   • 依赖控制与同步
   • 任务创建与执行流程
3. 关键函数与类
   • GPipe 类
   • fork 与 join 函数
   • record_stream 函数
   • depend 函数
   • fence 函数
   • compute 函数
   • worker 函数
   • spawn_workers 函数
   • Worker 线程通信
4. 跨设备同步实现
5. 总结

GPipe Pipeline 概述

GPipe 是一种管道并行（pipeline parallelism）实现，通常用于分布式深度学习中大模型的训练。它将模型划分为多个分区，每个分区放在不同的设备上，数据在设备之间逐步传递以完成计算。GPipe 使用检查点和数据依赖管理来优化显存利用并确保计算顺序。

核心模块解析

Fork 和 Join 的作用

fork 和 join 用于在计算图中建立伪依赖关系，确保反向传播的顺序。具体来说：

• fork：在前向传播中创建分叉，生成一个虚拟张量 phony，用于在计算图中保持依赖。
• join：将输出张量与 phony 合并，确保反向传播时遵循指定的顺序。

这种机制保证了反向传播的顺序正确，特别适合在多设备或多分区的流水线并行中使用。

依赖控制与同步

• fork 和 join 将分区的计算按顺序排布，避免资源竞争。
• 在反向传播时，join 通过 phony 依赖，确保在 B 之前计算 A 的梯度，从而实现同步。