hfai.pl | 兼具萤火集群优化特性的 PyTorch Lightning

 hfai.pl

Pytorch Lightning（简称 pl） 是在 PyTorch 基础上进行封装的库，它能帮助开发者脱离 PyTorch 一些繁琐的细节，专注于核心代码的构建，在 PyTorch 社区中备受欢迎。hfai.pl 是 high-flyer 对 pl 的进一步封装，能更加轻松的适配各种集群特性，带来更好的使用体验。本文将为大家详细介绍优化细节。

---

集群环境适配

若集群每个计算节点有 x 张 GPU。用户提交任务时需选定节点数量 N，则该任务可获得 N*x 个 GPU。每个进程中全局环境变量含义如下：

• world_size (hfai): 节点数量，用 N 表示

• rank (hfai): 节点 id，用 n 表示，n 属于 0 ~ N-1

• local_rank (hfai): GPU id, 用 k 表示，k 属于 0 ~ x-1

与 PyTorch init_process_group 的变量有如下对应：

• world_size (PyTorch): 进程数量，每个 GPU 对应一个进程，因此总进程数目为总 GPU 数目，计算方法为 N*x

• rank (PyTorch): 进程 id，利用节点 id 和 GPU id 可以计算出进程的 id，计算方法为 n*x+k

在一般的 PyTorch 代码中，我们需要进行如下的分布式初始化：

ip = os.environ.get("MASTER_ADDR", "127.0.0.1")
port = os.environ.get("MASTER_PORT", "2223")
hosts = int(os.environ.get("WORLD_SIZE", 1))  # number of nodes
rank = int(os.environ.get("RANK", 0))  # node id
gpus = torc