多 GPU JAX 训练如何开启 PGLE，让集合通信与计算重叠-开发者社区

多 GPU JAX 训练如何开启 PGLE，让集合通信与计算重叠

【免费下载链接】jaxComposable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/GitHub_Trending/ja/jax

在多 GPU 上跑 JAX 训练时，AllReduce、AllGather、ReduceScatter 这类集合通信操作如果与计算串行执行，GPU 时间就会被白白消耗在等待通信上。JAX 的 GPU performance tips 给出了一个官方方案：Profile Guided Latency Estimator（PGLE）。它先实测各计算指令和集合通信的真实运行时间，再把这份 profile 信息喂回 XLA 编译器，用于生成更好的调度决策——由--xla_gpu_enable_latency_hiding_scheduler延迟隐藏调度器把异步通信与计算重叠起来。该文档明确说明：与 GPU 间通信相关的优化只有在多设备运行时才有意义。

PGLE 支持 GPU 平台，有两种用法：

Auto PGLE：一次运行内自动完成「采集 profile → 重新编译」，对训练脚本最省事；
Manual PGLE：跑两次任务，第一次采集并保存 profile，第二次把 profile 文件喂给编译器。

一个必须先知道的前提：JAX profiler（两种 PGLE 工作流都依赖它）不能与 NVIDIA Nsight Systems 共存，这一点本文后文会给出规避方式。

开启 Auto PGLE（推荐主路径）

方式一：设置环境变量

在训练入口的 shell 环境或启动脚本中设置。必选：

export JAX_ENABLE_PGLE=true # For JAX version <= 0.5.0 make sure to include: export XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true"

注意第二条 XLA flag：文档标注它针对 JAX <= 0.5.0，如果你的 JAX 版本更新，是否仍需要以当前版本的默认行为为准，可参考 xla_flags 核对。

可选参数：

export JAX_PGLE_PROFILING_RUNS=3 export JAX_PGLE_AGGREGATION_PERCENTILE=85

JAX_PGLE_PROFILING_RUNS：控制用多少次重跑来采集 profile 数据。次数越多 profile 信息越好，但会多出等量的非优化训练步骤。默认值是 3（见 config.py 中pgle_profiling_runs的定义）。
JAX_PGLE_AGGREGATION_PERCENTILE：聚合各设备性能数据所用的百分位。当各 step 间性能抖动太大、噪声难以过滤时，可以降低该值。默认值是 90。

另外文档提示：目前 Auto PGLE 的 profile 采集与 command buffer 不兼容。若启用了 command buffer，Auto PGLE 会在采集期间临时禁用它、重编译后再恢复；如果你希望开启/关闭 PGLE 时 command buffer 行为一致，可以手动禁用：

export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_enable_command_buffer=''"

方式二：在 Python 代码中用 config 上下文设置

如果不想改 shell 环境，也可以在代码里用 JAX 的 config 上下文管理器（见 gpu_performance_tips.md）：

import jax from jax._src import config with config.enable_pgle(True), config.pgle_profiling_runs(1): # Run with the profiler collecting performance information. train_step() # Automatically re-compile with PGLE profile results train_step() ...

这里train_step()是你的训练步骤函数，需要替换为自己的实现。文档描述的行为是：在pgle_profiling_runs指定的次数内运行并采集性能信息，之后 JAX 自动用采集到的 profile 结果重新编译模块，后续调用直接走重编译后的可执行体。

限制：Auto PGLE 不适用于预编译模块

Auto PGLE 需要在运行时重新编译模块，因此对以下写法不生效：

AOT（提前编译）的模块；
先lower().compile()再调用的模式：

import jax from jax._src import config train_step_compiled = train_step().lower().compile() with config.enable_pgle(True), config.pgle_profiling_runs(1): train_step_compiled() # No effect since module was pre-compiled. train_step_compiled()

如果你使用持久编译缓存 + AOT 的流程，Auto PGLE 的开法不成立，应转看下文「配合持久编译缓存采集 Nsight 分析」的两步法。

验证与排查

判断 Auto PGLE 是否在正常工作：文档描述的预期行为是模块在采集够 profile 运行次数后于同一次运行内自动重编译，你不需要手动干预。tests/pgle_test.py 中的testAutoPgle展示了官方验证思路：开启xla_dump_to与xla_gpu_experimental_dump_fdo_profiles编译选项后，第 1 次运行生成的.fdo_profile文件应为空；第 3 次运行（触发重编译）之后会多出非空的 FDO profile 文件，说明 profile 被实际喂入了重编译。

PGLE 采到空 trace 时：当检测到 trace 为空，JAX 会发出RuntimeWarning，提示可能是与 Nsight Systems 等同样订阅 CUPTI 的工具产生了竞争（对应 XLA 侧的CUPTI_ERROR_MULTIPLE_SUBSCRIBERS_NOT_SUPPORTED），并建议改为「先用 PGLE 填充持久编译缓存，再用JAX_COMPILATION_CACHE_EXPECT_PGLE跑第二次」的方式（见 profiler.py 中的告警文案）。

需要 Nsight Systems 分析时：由于 JAX profiler 与 Nsight 不能共存，文档给出两步流程（JAX_COMPILATION_CACHE_EXPECT_PGLE配置项自 JAX v0.5.1 引入）：

第一步，启用 PGLE 把优化后的函数写进持久编译缓存：

export JAX_ENABLE_COMPILATION_CACHE=yes # not strictly needed, on by default export JAX_COMPILATION_CACHE_DIR=/root/jax_cache JAX_ENABLE_PGLE=yes python my-model.py

第二步，用 Nsight 跑第二次，并从缓存加载 PGLE 优化后的函数：

JAX_COMPILATION_CACHE_EXPECT_PGLE=yes nsys profile python my-model.py

my-model.py替换为你的训练入口；JAX_COMPILATION_CACHE_DIR示例值/root/jax_cache来自文档，可替换为你有写权限的目录。JAX_COMPILATION_CACHE_EXPECT_PGLE的语义是：编译缓存中若存在用 profile 数据编译过的条目，即使当前未启用 PGLE 也会优先加载；找不到优先条目时会打印告警。

Manual PGLE（可选替代路径）

如果 Auto PGLE 不适合你的流程（例如想在采集与优化之间做其他处理），文档给出了手动三步流程。

第 1 步：开延迟隐藏调度器跑一次负载：

export XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true"

第 2 步：用 JAX profiler 采集并后处理 profile，把提取出的指令延迟保存为二进制 protobuf 文件：

import os from etils import epath import jax from jax.experimental import profiler as exp_profiler # Define your profile directory profile_dir = 'gs://my_bucket/profile' jax.profiler.start_trace(profile_dir) # run your workflow # for i in range(10): # train_step() # Stop trace jax.profiler.stop_trace() profile_dir = epath.Path(profile_dir) directories = profile_dir.glob('plugins/profile/*/') directories = [d for d in directories if d.is_dir()] rundir = directories[-1] logging.info('rundir: %s', rundir) # Post process the profile fdo_profile = exp_profiler.get_profiled_instructions_proto(os.fspath(rundir)) # Save the profile proto to a file. dump_dir = rundir / 'profile.pb' dump_dir.parent.mkdir(parents=True, exist_ok=True) dump_dir.write_bytes(fdo_profile)

示例中的profile_dir = 'gs://my_bucket/profile'是文档给出的 GCS 路径，替换为你可写的 profile 目录。运行你的工作负载的那几行（示例中被注释的train_step()循环）需要自己打开。完成后，rundir（代码里会打印出来）下会生成profile.pb。

第 3 步：把 profile 文件通过--xla_gpu_pgle_profile_file_or_directory_path传给第二次编译：

export XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_pgle_profile_file_or_directory_path=/path/to/profile/profile.pb"

把/path/to/profile/profile.pb替换为第 2 步实际生成的文件路径。

验证方式：把日志级别调到包含INFO，跑真实工作负载：

export TF_CPP_MIN_LOG_LEVEL=0

如果运行日志里出现下面这两条（文档示例日志）：

2023-07-21 16:09:43.551600: I external/xla/xla/service/gpu/gpu_hlo_schedule.cc:478] Using PGLE profile from /tmp/profile/plugins/profile/2023_07_20_18_29_30/profile.pb 2023-07-21 16:09:43.551741: I external/xla/xla/service/gpu/gpu_hlo_schedule.cc:573] Found profile, using profile guided latency estimator

说明 profile 已被延迟隐藏调度器使用。

参考

主文档：docs/gpu_performance_tips.md（PGLE 自动/手动工作流、flag 说明、Nsight 两步法）
配置项定义与默认值：jax/_src/config.py（enable_pgle、pgle_profiling_runs默认 3、pgle_aggregation_percentile默认 90、compilation_cache_expect_pgle）
Auto PGLE 行为验证：tests/pgle_test.py、多机测试 tests/multiprocess/pgle_test.py

【免费下载链接】jaxComposable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/GitHub_Trending/ja/jax

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考