使用强化学习进行目标无关的 XLA 优化
机器学习
2023-08-29 v1 人工智能
摘要
在像 XLA 这样的机器学习编译器中的一个重大挑战是多遍优化与分析。近来主要关注于 XLA 在图级、子图级和核级阶段的目标相关优化。我们特别关注目标无关的 XLA HLO 遍排序优化：我们的方法旨在寻找编译器优化遍的最优序列，这与目标相关优化解耦。然而，在 XLA HLO 的遍排序方面鲜有特定领域的研究。为此，我们提出引入基于深度强化学习（RL）的搜索以寻找最优 XLA HLO 遍排序。我们还提出对深度 RL 算法的增强，以进一步提升最优搜索性能，并为 RL 的特定领域引导开辟研究方向。我们创建了 XLA Gym 实验框架作为工具，使 RL 算法能够与编译器交互以进行遍优化从而训练智能体。总体而言，在我们的实验中，观察到在 GPT-2 训练图基准上操作数减少平均提升 ，在包含 GPT-2、BERT 和 ResNet 图的多样化基准上较编译器默认阶段排序使用所提方法提升 。
引用
@article{arxiv.2308.14364,
title = {Target-independent XLA optimization using Reinforcement Learning},
author = {Milan Ganai and Haichen Li and Theodore Enns and Yida Wang and Randy Huang},
journal= {arXiv preprint arXiv:2308.14364},
year = {2023}
}
备注
Workshop on ML for Systems @ NeurIPS 2022