位元組跳動開源分散式訓練框架BytePS，支持PyTorch、TensorFlow等

機器之心報道

參與：戴一鳴、思源

最近，位元組跳動發布了一款通用高性能分散式訓練框架 BytePS，該框架支持TensorFlow、Keras、PyTorch 和 MXNet，並且可以在 TCP 或 RDMA 網路上運行。根據該項目的 GitHub 頁面，BytePS顯著優於目前的開源分散式訓練框架。例如在流行的公有雲和同樣數量 GPU 上，BytePS 的訓練速度可以達到Horovod (NCCL) 的兩倍。

最近，位元組跳動發布了一款通用高性能分散式訓練框架 BytePS，該框架支持 TensorFlow、Keras、PyTorch 和 MXNet，並且可以在 TCP 或 RDMA 網路上運行。

BytePS GitHub 地址：https://github.com/bytedance/byteps

根據該項目的 GitHub 頁面，BytePS 顯著優於目前的開源分散式訓練框架。例如在流行的公有雲和同樣數量 GPU 上，BytePS 的訓練速度可以達到 Horovod (NCCL) 的兩倍。

框架性能

開發團隊在 BytePS 上測試了兩個模型：VGG16（通信密集）和 Resnet50（計算密集）。測試使用了 Tesla V100 16GB GPU 集群，批大小都是 64。機器使用的是公有雲上的虛擬機，每個機器有 8 個 GPU，集成了 NVLink。機器之間使用 20 Gbps TCP/IP 網路互通。在測試上，BytePS 在 Resnet50 的表現較 Horovod（NCCL）提高 44%，在 VGG16 則提升了 100%。

在 VGG16 和 Resnet50 模型上，BytePS 和 NCCL 的性能對比。

作者提供了 Docker 文件，幫助用戶測試。

拋棄 MPI，迎接雲計算

為什麼 BytePS 的表現比 Horovod 好太多？主要原因是 BytePS 為雲計算和共享集群設計，並拋棄了 MPI。

MPI 是高性能計算（High Performance Computing）的產物。當需要運行單一任務時，MPI 適合用於建立在同質化硬體的計算集群上。但是雲計算（或者內部共享集群）是不一樣的。

團隊因此重新思考了最佳通信策略。總之，BytePS 不僅在機器內使用 NCCL，同時也重新部署了機器間的通信方式。

BytePS 同時繼承了許多加速技術，如分級策略、管道、張量分割、NUMA-aware 本地通信、基於優先順序的調度機制等。

快速上手

使用 BytePS 前，假設你已經安裝了以下一種或更多框架：TensorFlow、Keras、PyTorch、MXNet 等。BytePS 基於 CUDA 和 NCCL。

複製 BytePS 和第三方依賴：

git clone --recurse-submodules https://github.com/bytedance/byteps

然後進入 BytePS 文件目錄，並安裝：

python setup.py install

注意：你可能需要設置 BYTEPS_USE_RDMA=1 來安裝 RDMA 支持。

現在你可以試試一些例子。假設你使用 MXNet，並想嘗試 ResNet50 的基本模型。

export NVIDIA_VISIBLE_DEVICES=0,1

DMLC_NUM_WORKER=1

DMLC_NUM_SERVER=1

DMLC_WORKER_ID=0

DMLC_ROLE=worker

DMLC_PS_ROOT_URI=10.0.0.1

DMLC_PS_ROOT_PORT=1234

DMLC_INTERFACE=eth0

python byteps/launcher/launch.py byteps/example/mxnet/train_imagenet_byteps.py --benchmark 1 --batch-size=32

對於分散式訓練，你可能需要建立一個伺服器鏡像。團隊提供了 Docker 文件作為例子。你可以使用同樣的鏡像用於調度和提供服務。

更多啟動分散式任務和上手教程可參考：https://github.com/bytedance/byteps/tree/master/docs

在你的代碼中使用 BytePS

BytePS 和 Horovod 介面高度兼容。選擇 Horovod 可以降低測試工作量。

如果你的任務只依賴 Horovod 的 allreduce 和廣播，你可以在一分鐘內切換到 BytePS。

只需要用 import byteps.tensorflow as bps 替換 import horovod.tensorflow as hvd，並將代碼中所有的 hvd 替換成 bps。

項目計劃

BytePS 目前不支持單純的 CPU 訓練，一些底層邏輯可能不支持。你可以使用 CUDA 或 NCCL 來運行 BytePS。

BytePS 即將增加以下特性：

• 非同步訓練
• 容錯機制
• 延遲減緩

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！