XLA GPU 如何配置持久化 autotuning 缓存加速 Triton 内核编译？

XLA GPU 如何配置持久化 autotuning 缓存加速 Triton 内核编译【免费下载链接】tensorflowAn Open Source Machine Learning Framework for Everyone项目地址: https://gitcode.com/GitHub_Trending/te/tensorflow在 GPU 上用 XLA 编译模型时部分 GPU 内核由 OpenAI Triton 生成。Triton 能对特定 fusion 生成较快的内核但 XLA 需要为每一种 fusion 调优参数fusion 数量多时这个 autotuning 过程会显著拉长编译时间。XLA 的 Persisted autotuning 功能仅支持 GPU可以把这些调优结果持久化下来命中缓存的 fusion 直接复用结果未命中的 fusion 照常 autotune其余编译步骤不受影响。本文基于 third_party/xla/docs/persisted_autotuning.md 说明如何配置缓存、如何在测试中使用以及缓存何时会失效。前提该功能的适用范围仅适用于 GPU文档标题即为 Persisted autotuning (GPU only)。该功能默认关闭不传对应参数时缓存不生效。缓存中已有 fusion 的结果会被复用新 fusion 仍会正常 autotune因此缓存不会让首次遇到 fusion 的运行变快而是让后续遇到相同 fusion 的运行变快。推荐路径使用按 fusion 划分的缓存目录官方推荐的方式是通过参数指定一个缓存目录--xla_gpu_per_fusion_autotune_cache_diryour/directory其中your/directory替换为你实际使用的目录路径。启用后XLA 会在该目录中维护按 fusion 划分的 autotune 缓存每种不同的 fusion 对应一个文件。XLA 在需要时读取已有结果在结果确定后写入新结果。这种方式的两个好处同一个缓存目录可以被多次 XLA 运行甚至不同模型复用每遇到一个新 fusion 缓存就增长一份后续运行持续受益对多个 XLA 实例并发使用同一缓存目录提供了基本支持。启用前必须满足以下条件这些条件需要由使用方自行保证目录必须事先存在且可写——在运行 XLA 之前创建好目录并确认可写。缓存失效由用户处理——如果想从空缓存开始请提供一个空目录。XLA 版本检查由用户处理——如果想为不同版本的 XLA 使用独立缓存请使用不同的目录。关于缓存是否生效文档给出的行为描述是XLA 在需要时读取已有结果在新结果确定后写入。因此可以核对目录内容来确认缓存正在积累——运行后目录中应出现对应的 fusion 缓存文件且后续运行时已有 fusion 直接走缓存、其余 fusion 正常 autotune从而缩短后续编译时间。替代路径将整份 HLO 的调优结果写入单个文件除了目录方式autotuning 结果也可以从单个文件加载或导出--xla_gpu_dump_autotune_results_to --xla_gpu_load_autotune_results_from指定.txt或.textproto文件时缓存以 textproto 格式写入/读取其他扩展名则以二进制 protobuf 格式处理。注意限制目录方式与文件方式不保证能很好地组合使用。两种方式应二选一不要在同一次运行中混用。在 Bazel 测试中使用持久化 autotuning如果测试规模很大、尤其是测试环境性能有限时文档建议使用持久化 autotuning。需要注意缓存只有在与测试运行相同类型的 GPU上生成时才能正确工作若运行测试的 GPU 与缓存生成时的 GPU 不同缓存会被加载但不会被使用这一点也见于 tests/BUILD 中测试目标的注释。以一个固定使用某类 GPU 的测试为例流程分三步第 1 步从测试中导出 autotune 结果给测试命令追加以下参数TEST_UNDECLARED_OUTPUTS_DIR是 Bazel 测试的输出目录变量--test_envXLA_FLAGS--xla_gpu_dump_autotune_results_toTEST_UNDECLARED_OUTPUTS_DIR/autotune_cache.textproto --test_sharding_strategydisabled导出缓存时必须禁用 sharding--test_sharding_strategydisabled才能为所有测试正确得到单一 autotune 缓存。测试运行后autotune_cache.textproto会出现在测试的输出目录中。第 2 步把导出的缓存文件放入代码仓库供后续测试作为数据依赖使用。第 3 步把缓存加入测试目标的数据依赖并通过环境变量加载data [test_autotune_cache.textproto], env {XLA_FLAGS: --xla_gpu_load_autotune_results_from $(execpath test_autotune_cache.textproto)},与导出不同加载 autotune 结果的测试允许使用 sharding。上述流程在 third_party/xla/xla/backends/gpu/tests/BUILD 中有三个可直接参考的示例目标load_autotune_results_using_execpath_test通过$(execpath ...)加载缓存load_autotune_results_from_test_workspace_test通过TEST_WORKSPACE/前缀加载文档同时给出等价的手动命令行形式--test_envXLA_FLAGS--xla_gpu_load_autotune_results_fromTEST_WORKSPACE/my/package/autotune_results_test.textprotodump_autotune_results_to_test_outputs_test把缓存导出到TEST_UNDECLARED_OUTPUTS_DIR该目标设置了shard_count 1配合禁用的 sharding。示例中用到的数据文件是 test_autotune_cache.textproto。缓存何时失效以及如何再生成文档列出了三种缓存过期cache obsolescence的情形出现后需要重新生成 autotuning 缓存模型发生较多改动——缓存中可能不再包含全部 fusion测试会变慢此时应重新生成缓存开始使用新类型的 GPU 运行测试——缓存只对相同 GPU 类型有效XLA 编译器演进生成了不同的 fusion——缓存中的 fusion 集合可能不再匹配。对应的处理方式是按前面导出缓存的流程重新 dump 一份新缓存替换旧文件目录方式则按需要更换目录或清空后重建。限制与边界该功能仅限 GPU且默认关闭必须显式传入参数才生效。目录缓存要求目录预先存在且可写缓存失效与 XLA 版本区分都由使用方负责。目录方式与单文件 dump/load 方式不保证能良好组合选择其中一种即可。缓存不会让第一次遇到 fusion 的运行加速收益体现在后续运行中已有 fusion 走缓存新 fusion 正常 autotune。【免费下载链接】tensorflowAn Open Source Machine Learning Framework for Everyone项目地址: https://gitcode.com/GitHub_Trending/te/tensorflow创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考