注意
Ray 2.40 默认使用 RLlib 的新 API 栈。Ray 团队已基本完成将算法、示例脚本和文档迁移到新的代码库。
如果你仍在使旧的 API 栈，请参阅新 API 栈迁移指南了解如何迁移。
将 RLlib 与 torch 2.x compile 结合使用#
torch 2.x 附带 torch.compile() API，可用于 JIT 编译包装的代码。我们在 RL 模块 和 学习器 的上下文中将 torch.compile() 集成到 RLlib 中。
我们将此功能与 RL 模块集成了。你可以在 AlgorithmConfig 对象上通过 framework() API 设置后端和模式。另外，你也可以在独立使用时（例如推理）直接编译 RLModule。
基准测试#
我们对此功能进行了全面基准测试。以下基准测试仅考虑在推理和环境探索期间启用 torch-compile 可能带来的加速效果。这种加速方法很重要，因为强化学习通常受采样瓶颈的限制。
推理#
对于基准测试指标，我们计算运行 RLModule 的 forward_exploration() 所花费时间的倒数。我们在不同硬件设置、torch 版本、dynamo 后端和模式以及不同批量大小下，对 PPO RLModule 的默认实现进行了此基准测试。下表显示了在给定硬件和 PyTorch 版本组合下，我们发现的能产生最高加速效果的 torch-backend 和 -mode 组合
硬件 |
PyTorch 版本 |
加速 (%) |
后端 + 模式 |
|---|---|---|---|
CPU |
2.0.1 |
33.92 |
ipex + default |
CPU |
2.1.0 nightly |
x |
ipex + default |
T4 |
2.0.1 |
14.05 |
inductor + reduce-overhead |
T4 |
2.1.0 nightly |
15.01 |
inductor + reduce-overhead |
V100 |
2.0.1 |
92.43 |
inductor + reduce-overhead |
V100 |
2.1.0 nightly |
85.71 |
inductor + reduce-overhead |
A100 |
2.0.1 |
x |
inductor + reduce-overhead |
A100 |
2.1.0 nightly |
156.66 |
inductor + reduce-overhead |
有关详细表格，请参阅附录。有关基准测试代码，请参阅run_inference_bm.py。要运行基准测试，请使用以下命令
python ./run_inference_bm.py --backend <dynamo_backend> --mode <dynamo_mode> -bs <batch_size>
一些元级评论#
性能提升取决于许多因素，包括使用的神经网络架构、采样时的批量大小、后端、模式、torch 版本以及许多其他因素。为了优化性能，首先让非编译工作负载学习，然后在不同硬件上对 torch compile 参数进行超参数调优。
对于 CPU 推理，请使用推荐的仅推理后端：
ipex和onnxrt。在 A100 等较新架构上，加速效果比 T4 等较旧架构更显著。
Torch compile 仍在发展中。我们注意到 2.0.1 版本和 2.1 nightly 版本之间存在显著差异。因此，在对自己的工作负载进行基准测试时，考虑 torch 版本非常重要。
探索#
在 RLlib 中，你现在可以设置配置，使其在 RL 智能体训练过程的采样阶段使用编译后的模块。默认情况下，rollout worker 在 CPU 上运行，因此建议使用 ipex 或 onnxrt 后端。但是，你也可以通过设置 num_gpus_per_env_runner 在 GPU 上运行采样部分，在这种情况下也可以使用其他后端。要在训练期间启用 torch-compile，你还可以设置 torch_compile_learner 等效项。
from ray.rllib.algorithms.ppo import PPOConfig
config = PPOConfig().framework(
"torch",
torch_compile_worker=True,
torch_compile_worker_dynamo_backend="ipex",
torch_compile_worker_dynamo_mode="default",
)
此基准测试脚本使用 Atari-Breakout 游戏的默认模型架构运行 PPO 算法。它对编译和未编译的 RLModules 进行 n 次迭代训练，并报告加速效果。请注意，负加速值意味着编译模块后会变慢。
要运行基准测试脚本，你需要一个 Ray 集群，至少包含 129 个 CPU (2x64 + 1) 和 2 个 GPU。如果你无法访问此配置，可以更改采样 worker 的数量和批量大小来降低要求。
python ./run_ppo_with_inference_bm.py --backend <backend> --mode <mode>
以下是结果摘要
后端 |
模式 |
加速 (%) |
|---|---|---|
onnxrt |
default |
-72.34 |
onnxrt |
reduce-overhead |
-72.72 |
ipex |
default |
11.71 |
ipex |
reduce-overhead |
11.31 |
ipex |
max-autotune |
12.88 |
如你所见，在我们测试的设置中，onnxrt 没有获得任何加速（实际上它使工作负载减慢了 70%），而 ipex 提供了约 10% 的加速。如果我们改变模型架构，这些数字可能会改变。因此，首先确定架构，然后再寻找最快的训练设置非常重要。