haiscale | 幻方萤火高性能并行训练工具库

最新推荐文章于 2024-11-21 10:12:10 发布

幻方AI小编

最新推荐文章于 2024-11-21 10:12:10 发布

阅读量495

点赞数

分类专栏：幻方AI训练平台 hfai深度学习套件幻方AI萤火文章标签：深度学习人工智能数据仓库 python

本文链接：https://blog.youkuaiyun.com/weixin_66945478/article/details/128287094

版权

haiscale是一个轻量级的高性能并行训练库，整合了幻方AI的并行训练经验，支持PyTorch用户在大规模集群上高效训练模型。主要特性包括：haiscale.ddp（分布式数据并行，基于hfreduce通信）、haiscale.fsdp（优化的FSDP实现）、haiscale.pipeline（流水线并行工具，支持GPipe、PipeDream等）和haiscale.cpu_offload（神经网络模型Offload）。实验证明，haiscale在各种并行方式下相比PyTorch官方工具有显著性能提升。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

haiscale (Highflyer AI Scale) 是一个轻量级的高性能并行训练工具库，其整合了幻方 AI 多年的并行训练研发优化经验，能够帮助 PyTorch 用户更加高效、便捷地在大规模集群上训练模型。

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

1. haiscale.ddp: 分布式数据并行工具，以幻方 AI 自研的 hfreduce 通信为后端，相比于 NCCL 能够获得更好的多卡拓展性能；

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

3. haiscale.pipeline: 分布式流水线并行（或称模型并行）工具包，包含 GPipe, PipeDream 等算法，支持多机多卡训练；

4. haiscale.cpu_offload: 神经网络模型 Offload 工具，节省训练占用的显存。

下图展示了 haiscale 三种并行方式的性能，其相比 PyTorch 官方自带工具都有显著的性能提升：

用于测试的模型是 GPT-2 Medium，相关代码已开源至 hfai 模型仓库。下面将为大家简要介绍。

API文档：https://doc.hfai.high-flyer.cn/api/haiscale_ddp.html

示例模型：https://github.com/HFAiLab/hfai-models/tree/main/gpt

安装

haiscale 提供 Python 接口，通过如下方式安装：

如果要使用 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