HaiScale 简介

haiscale (Highflyer AI Scale) 是一个轻量级的高性能并行训练工具库，能够帮助 PyTorch 用户更加高效、便捷地在大规模集群上训练模型。

haiscale 中包含了以下几种工具：

1. haiscale.ddp: 分布式数据并行工具，以 hfreduce 为后端，相比于 NCCL 能够获得更好的拓展性能

2. haiscale.fsdp: 极致优化的 Fully Sharded Data Parallel 实现，相比于 PyTorch FSDP 速度更快、占用显存更少

3. haiscale.pipeline: 分布式流水线并行工具包，包含 GPipe, PipeDream 等算法，支持多机训练

4. haiscale.cpu_offload: activation offload 工具，节省训练占用的显存

Performance

我们在 GPT-2 Medium 模型上测试三种不同的并行方式的性能，并和 PyTorch 官方自带的工具做比较，以下是运行的结果：

Installation

1. 如果要使用 haiscale DDP，我们需要先安装 hfreduce （如果不需要使用 DDP 可跳过这步）：

   sudo apt install libnuma-dev
   sudo apt install libibverbs-dev
   pip install hfreduce --extra-index-url https://pypi.hfai.high-flyer.cn/simple --trusted-host pypi.hfai.high-flyer.cn
   

2. 安装 haiscale：

   pip install haiscale --extra-index-url https://pypi.hfai.high-flyer.cn/simple --trusted-host pypi.hfai.high-flyer.cn
   

Usage

使用方法见后续 API 文档

Examples

我们提供了一个 GPT-2 Medium 的训练示例，请到 模型仓库 下查看