备注
Ray 2.10.0 引入了 RLlib 的“新 API 栈”的 alpha 阶段。Ray 团队计划将算法、示例脚本和文档迁移到新的代码库中，从而在 Ray 3.0 之前的后续小版本中逐步替换“旧 API 栈”（例如，ModelV2、Policy、RolloutWorker）。
然而，请注意，到目前为止，只有 PPO（单代理和多代理）和 SAC（仅单代理）支持“新 API 堆栈”，并且默认情况下继续使用旧 API 运行。您可以继续使用现有的自定义（旧堆栈）类。
请参阅此处 以获取有关如何使用新API堆栈的更多详细信息。
使用 RLlib 与 torch 2.x 编译#
torch 2.x 带来了 torch.compile() API，它可以用于即时编译（JIT-compile）包装的代码。我们在 RLModules 和 Learners 的上下文中将 torch.compile() 与 RLlib 集成。
我们已经将此功能与 RLModules 集成。您可以通过 AlgorithmConfig 对象上的 framework() API 设置后端和模式。或者，您可以直接在独立使用期间（如推理）编译 RLModule。
基准测试#
我们对此功能进行了全面的基准测试。以下基准测试仅考虑在推理和环境探索期间启用 torch-compile 可能带来的潜在速度提升。这种方法是相关的，因为强化学习通常受采样瓶颈的限制。
推理#
对于基准测试指标，我们计算运行 forward_exploration() 方法所需时间的倒数。我们在不同的硬件设置、torch 版本、dynamo 后端和模式以及不同的批量大小下，对 PPO RLModule 的默认实现进行了此基准测试。下表显示了在给定硬件和 PyTorch 版本组合下，能够获得最高加速的 torch-backend 和 -mode 的组合：
硬件 |
PyTorch 版本 |
加速（%） |
后端 + 模式 |
|---|---|---|---|
CPU |
2.0.1 |
33.92 |
ipex + 默认 |
CPU |
2.1.0 夜间版 |
x |
ipex + 默认 |
T4 |
2.0.1 |
14.05 |
电感器 + 减少开销 |
T4 |
2.1.0 夜间版 |
15.01 |
电感器 + 减少开销 |
V100 |
2.0.1 |
92.43 |
电感器 + 减少开销 |
V100 |
2.1.0 夜间版 |
85.71 |
电感器 + 减少开销 |
A100 |
2.0.1 |
x |
电感器 + 减少开销 |
A100 |
2.1.0 夜间版 |
156.66 |
电感器 + 减少开销 |
有关详细表格，请参见 附录 。有关基准测试代码，请参见 run_inference_bm.py 。要运行基准测试，请使用以下命令：
python ./run_inference_bm.py --backend <dynamo_backend> --mode <dynamo_mode> -bs <batch_size>
一些元级别的评论#
性能提升取决于许多因素，包括使用的神经网络架构、采样时的批量大小、后端、模式、torch版本以及其他许多因素。为了优化性能，首先获取非编译工作负载的学习，然后在不同硬件上对torch编译参数进行超参数调优。
对于CPU推理，请使用推荐的仅推理后端：
ipex和onnxrt。与旧款如T4相比，在更现代的架构如A100上，加速效果更为显著。
Torch 编译仍在不断发展。我们注意到 2.0.1 版本和 2.1 夜间版本之间存在显著差异。因此，在基准测试您自己的工作负载时，考虑 torch 版本是非常重要的。
探索#
在 RLlib 中，您现在可以设置配置，以便在 RL 代理训练过程的采样期间使用编译模块。默认情况下，回滚工作程序在 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",
)
This 基准脚本在Atari-Breakout游戏中使用默认模型架构运行PPO算法。它对编译和未编译的RLModules进行``n``次迭代训练，并报告加速比。请注意，负的加速比值意味着编译模块时会减慢速度。
要运行基准测试脚本，您需要一个由至少129个CPU（2x64 + 1）和2个GPU组成的Ray集群。如果此配置对您不可用，您可以更改采样工作者的数量和批量大小以减少要求。
python ./run_ppo_with_inference_bm.py --backend <backend> --mode <mode>
以下是结果的总结：
后端 |
模式 |
加速（%） |
|---|---|---|
onnxrt |
默认 |
-72.34 |
onnxrt |
reduce-overhead |
-72.72 |
ipex |
默认 |
11.71 |
ipex |
reduce-overhead |
11.31 |
ipex |
max-autotune |
12.88 |
如你所见，onnxrt 在我们测试的设置中并没有获得任何速度提升（实际上它使工作负载减慢了70%），而 ipex 提供了约10%的速度提升。如果我们改变模型架构，这些数字可能会发生变化。因此，首先固定架构，然后寻找最快的训练设置是非常重要的。