Ray Train 基准测试#

下面我们记录了常见 Ray Train 任务和工作流的关键性能基准测试。

GPU 图像训练#

此任务使用 TorchTrainer 模块，利用 PyTorch ResNet 模型训练不同大小的数据。

我们测试了不同集群大小和数据大小下的性能。

注意

对于多主机分布式训练，在 AWS 上，我们需要确保 EC2 实例位于同一个 VPC 中，并且安全组中所有端口都已打开。

集群设置

数据大小

性能

命令

1 个 g3.8xlarge 节点（1 个 worker）

1 GB (1623 张图片)

79.76 秒（2 个 epoch，40.7 张图片/秒）

python pytorch_training_e2e.py --data-size-gb=1

1 个 g3.8xlarge 节点（1 个 worker）

20 GB (32460 张图片)

1388.33 秒（2 个 epoch，46.76 张图片/秒）

python pytorch_training_e2e.py --data-size-gb=20

4 个 g3.16xlarge 节点（16 个 worker）

100 GB (162300 张图片)

434.95 秒（2 个 epoch，746.29 张图片/秒）

python pytorch_training_e2e.py --data-size-gb=100 --num-workers=16

PyTorch 训练性能对比#

此任务检查原生 PyTorch 分布式训练与 Ray Train 分布式 TorchTrainer 之间的性能对比。

我们证明了这两种框架的性能相似（差异在 2.5% 以内）。性能在不同的模型、硬件和集群配置下可能差异很大。

报告的时间为原始训练时间。两种方法都有一个未报告的几秒固定设置开销，对于较长时间的训练运行来说可以忽略不计。

集群设置

数据集

性能

命令

4 个 m5.2xlarge 节点（4 个 worker）

FashionMNIST

196.64 秒（对比 PyTorch 194.90 秒）

python workloads/torch_benchmark.py run --num-runs 3 --num-epochs 20 --num-workers 4 --cpus-per-worker 8

4 个 m5.2xlarge 节点（16 个 worker）

FashionMNIST

430.88 秒（对比 PyTorch 475.97 秒）

python workloads/torch_benchmark.py run --num-runs 3 --num-epochs 20 --num-workers 16 --cpus-per-worker 2

4 个 g4dn.12xlarge 节点（16 个 worker）

FashionMNIST

149.80 秒（对比 PyTorch 146.46 秒）

python workloads/torch_benchmark.py run --num-runs 3 --num-epochs 20 --num-workers 16 --cpus-per-worker 4 --use-gpu

TensorFlow 训练性能对比#

此任务检查原生 TensorFlow 分布式训练与 Ray Train 分布式 TensorflowTrainer 之间的性能对比。

我们证明了这两种框架的性能相似（差异在 1% 以内）。性能在不同的模型、硬件和集群配置下可能差异很大。

报告的时间为原始训练时间。两种方法都有一个未报告的几秒固定设置开销，对于较长时间的训练运行来说可以忽略不计。

注意

GPU 基准测试的批处理大小和 epoch 数量不同，导致运行时间更长。

集群设置

数据集

性能

命令

4 个 m5.2xlarge 节点（4 个 worker）

FashionMNIST

78.81 秒（对比 TensorFlow 79.67 秒）

python workloads/tensorflow_benchmark.py run --num-runs 3 --num-epochs 20 --num-workers 4 --cpus-per-worker 8

4 个 m5.2xlarge 节点（16 个 worker）

FashionMNIST

64.57 秒（对比 TensorFlow 67.45 秒）

python workloads/tensorflow_benchmark.py run --num-runs 3 --num-epochs 20 --num-workers 16 --cpus-per-worker 2

4 个 g4dn.12xlarge 节点（16 个 worker）

FashionMNIST

465.16 秒（对比 TensorFlow 461.74 秒）

python workloads/tensorflow_benchmark.py run --num-runs 3 --num-epochs 200 --num-workers 16 --cpus-per-worker 4 --batch-size 64 --use-gpu

XGBoost 训练#

此任务使用 XGBoostTrainer 模块，通过不同程度的并行化在不同大小的数据上进行训练，以展示分布式数据并行化接近线性的扩展能力。

对于此任务，XGBoost 参数使用 xgboost==1.7.6 的默认值。

集群设置

分布式训练 worker 数量

数据大小

性能

命令

1 个配备 16 个 CPU 的 m5.4xlarge 节点

1 个训练 worker 使用 12 个 CPU，为 Ray Data 任务留下 4 个 CPU

10 GB（2600 万行）

310.22 秒

python train_batch_inference_benchmark.py "xgboost" --size=10GB

10 个 m5.4xlarge 节点

10 个训练 worker（每个节点一个），使用 10x12 个 CPU，为 Ray Data 任务留下 10x4 个 CPU

100 GB（2.6 亿行）

326.86 秒

python train_batch_inference_benchmark.py "xgboost" --size=100GB