由于 JAX with Pathways 的分布式特性，某些操作可能因通信开销而无法很好地扩缩。虽然 Pathways 通过异步调度等功能最大限度地减少了这些开销，但在将 JAX 工作负载移植到 Pathways 或将 JAX with Pathways 工作负载扩缩到大量加速器时，您需要注意一些事项。
准备工作
请确保您已备妥：
进程索引
JAX with Pathways 会将 Pathways 集群中的所有设备视为本地设备。这简化了设备管理，并允许 JAX 利用所有可用资源。在实践中，这意味着：
- 对于所有设备，
jax.process_index()始终为 0。 jax.devices()和jax.local_devices()会返回整个作业中的所有 TPU 设备。
硬件类型和主机托管
为获得最佳性能，请将所有 Pathways 组件和用户作业放在同一 Google Cloud 云可用区中。使用大型 CPU，例如 IFRT 代理和资源管理器。我们建议至少使用专用 n2-standard-64，它配备 64 个 vCPU 和 256 GB 内存。
PathwaysUtils
Pathways-utils 是一个基于 Python 的 GitHub 代码库，提供必要的实用程序和工具，可让您简化在 Pathways on Cloud 架构上部署和执行 JAX 工作负载的过程。此软件包可处理云环境所需的调整，让 JAX 开发者能够专注于其核心机器学习工作流，只需进行极少的平台特定配置。具体而言，它提供：
- “代理”JAX 后端：此自定义后端使您的 JAX 应用能够通过设置
JAX_PLATFORMS=proxy环境变量来使用 Pathways 基础设施。 - 集成式性能分析实用程序：性能分析功能，可让您了解应用的性能。通过使用
jax.profiler.start_trace和jax.profiler.start_server等标准 JAX 性能分析 API，您不仅可以对 JAX 代码进行性能分析，还可以对底层 Pathways 组件进行性能分析，从而全面了解云环境中的执行情况。 - 使用 Orbax 进行分布式检查点：一种自定义 Orbax 检查点处理程序，可让您在 Pathways 环境中使用 Orbax 库时使用分布式检查点并恢复检查点。此集成旨在无需对现有 Orbax 检查点代码进行任何更改即可运行，只要该代码导入
pathwaysutils即可。 - 弹性训练基元：提供基础弹性训练基元，您可以使用这些基元通过 Pathways 构建强大且可伸缩的训练工作流。这些基元可让您的训练作业动态适应可用资源的变化，从而提高云环境中的效率和弹性。
检查点
Orbax 已经过 Pathways 的全面测试，可用于通过 Cloud Storage 进行分布式检查点和恢复。当您
设置环境变量 ENABLE_PATHWAYS_PERSISTENCE=1 并在 train.py 中调用
import pathwaysutils; pathwaysutils.initialize() 时，系统会注册一个自定义
ArrayHandler，该
可通过 IFRT
代理高效处理检查点操作，从而允许加速器上的 Pathways 工作器直接保存和恢复数据。
共置 Python
Colocated Python 是一个开源 JAX API，可让您直接在 TPU 或 GPU 主机上运行用户指定的 Python 代码，这在多控制器 JAX 中更为直接。通过使用此方法，您可以避免在客户端和 TPU 机器之间传输数据，以执行计算密集型任务，例如数据加载和检查点。如需将 Pathways 集群配置为运行共置的 Python JAX API，请使用示例边车 Dockerfile 构建并推送边车容器映像。
数据加载
在训练期间，我们会反复从数据集中加载批量数据，以馈送给模型。为了避免加速器缺少工作，请务必使用高效的异步数据加载器，将批次分片到多个主机。使用 Pathways 运行训练时，数据加载器在 CPU 虚拟机（与多控制器设置中使用的 TPU 虚拟机不同）上运行，并将数据分派到 TPU 虚拟机。这会导致读取数据的延迟时间较长，但通过在 CPU 主机上预读 X 个批次并将读取的数据异步分派到 TPU，可以部分缓解这种情况。此解决方案足以满足中小型规模的运行需求。
在 GKE 上的头工作负载容器中运行多进程数据加载器（例如 Grain）时，该容器会与 pathways-proxy 和 pathways-rm 边车共享头节点。由于此共享架构，您必须配置以下内容：
- 设置明确的工作器数量：在容器内，
os.cpu_count()会报告主机虚拟机的总 vCPU 数量，而不是容器的 CPU 限制。请传递明确的工作器数量（不超过容器的 CPU 限制），而不是从os.cpu_count()派生。 - 装载内存支持的
/dev/shm卷：多进程数据加载器使用 POSIX 共享内存 (/dev/shm) 进行进程间通信。确保在/dev/shm处装载了内存支持的emptyDir(medium: Memory) 卷，以便预提取的批次不会超出默认容器共享内存限制。
为了在大规模运行时获得最佳性能，我们强烈建议您使用共置 Python 直接在加速器上运行数据流水线，从而共置输入数据流水线。这样可以消除 CPU 瓶颈，并利用 TPU 的快速互连进行数据传输。
您可以在 multihost_dataloading.py 的 RemoteIterator 实现中找到迁移基于 TFDS 的输入流水线的参考实现。此实现使用共置的 Python JAX API，以分布式方式在多控制器 JAX 和 Pathways 上运行。
Jax 版本控制
Pathways 版本与 JAX 版本紧密耦合，以确保兼容性和稳定性。为避免潜在问题，请验证您的 Pathways 工件和 JAX 版本是否一致。每个 Pathways 版本都会通过 jax-<version> 格式的标记明确指定兼容的 JAX 版本。
编译缓存
Pathways 永久性编译缓存是一项功能，可让 Pathways 服务器将编译后的 XLA 可执行文件存储在永久性位置（例如 Cloud Storage），以避免冗余编译。默认情况下，此功能处于启用状态。缓存的位置作为 --gcs_scratch_location 标志传递给资源管理器和 Pathways 工作器容器。为了将相关存储费用降至最低，缓存会将生命周期政策附加到 Cloud Storage 位置。每个 Cloud Storage 存储桶最多只能有 50 项政策。因此，我们建议在所有工作负载中使用一个通用的 Cloud Storage 位置。
如果您省略 --gcs_scratch_location 标志或将其留空，Pathways 会自动回退到内存中编译缓存。此缓存不会在 Pathways 部署中保留已编译的 XLA 可执行文件。
此缓存类似于 JAX 编译缓存，后者由 pathwaysutils.initialize() 针对 Pathways 工作负载停用。
编译缓存需要以下 Cloud Storage 权限：
storage.buckets.get：检索存储桶元数据。storage.buckets.update：对于 Pathways 设置对象生命周期政策以强制执行 TTL 进行缓存逐出至关重要。storage.objects.list：列出存储桶中的现有缓存对象。storage.objects.create：将新编译的可执行文件写入缓存。storage.objects.get：从存储桶中读取缓存的可执行文件。
分析
您可以使用 JAX 分析器生成 JAX 程序的跟踪记录。Pathways 支持两种常见方式：
在这两种情况下，分析结果都会写入 Cloud Storage 存储桶。 Cloud Storage 存储桶中可能会创建多个跟踪记录文件，这些文件可能位于不同的时间戳文件夹下，例如：
- 调用跟踪记录的主 Python 进程（通常是笔记本虚拟机）：
<jax-client-vm-name>.xplane.pb - Pathways IFRT 代理：
client.<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb - Pathways 资源管理器：
server.*<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb - Pathways 工作器：
server.*<tpu-node-name>.xplane.pb
您可以通过运行以下命令，使用 TensorBoard 分析这些跟踪文件。如需详细了解 TensorBoard 及其所有性能分析工具，请参阅使用 Profiler 优化 TensorFlow 性能。
# verify trace files are present gsutil ls -l -r gs://BUCKET/PREFIX # View on tensorboard tensorboard --logdir=gs://BUCKET/PREFIX
替换以下内容：
BUCKET：用于存储跟踪文件的 Cloud Storage 存储桶PREFIX：Cloud Storage 存储桶中用于存储跟踪文件的路径
程序化性能剖析文件捕获
从代码内部捕获性能剖析。这些配置文件保存在 gs://<bucket>/<prefix> 内的时间戳目录下
import jax import pathwaysutils pathwaysutils.initialize() jax.profiler.start_trace("gs://BUCKET/PREFIX") # Run the operations to be profiled key = jax.random.key(0) x = jax.random.normal(key, (5000, 5000)) y = x @ x y.block_until_ready() jax.profiler.stop_trace()
手动捕获性能剖析文件
如需手动捕获性能分析信息，您必须从 Python 代码启动性能分析器服务器：
import jax
import pathwaysutils
pathwaysutils.initialize()
jax.profiler.start_server(jax_profiler_port)
# Your JAX code
jax.profiler.stop_server() # this is functionally a no-op
在性能分析器服务器运行时，您可以捕获性能剖析文件并将数据导出到目标 Cloud Storage 位置：
export DURATION_IN_SECS=6
export OUTPUT_DIR=gs://<var>BUCKET</var>/<var>PREFIX</var>
curl -d "{\"duration_ms\":\"${DURATION_IN_SECS} * 1000 }}\", \"repository_path\":\"${OUTPUT_DIR}\"}" -H "Content-Type: application/json" -X POST http://localhost:<jax_profiler_port>
您可以在程序的跟踪记录中找到 IFRT 代理客户端方法（例如 Compile 和 Execute）的时间信息。这些事件详细说明了在编译和执行期间与 IFRT 代理 gRPC 服务器的交互，它们会显示在名为 GrpcClientSessionUserFuturesWorkQueue 的线程上。通过在跟踪记录中检查此线程，您可以深入了解这些操作的性能。
XLA 标志
使用 Pathways 时，您必须在 pathways-proxy 容器中设置 XLA 标志。
使用 Cluster Toolkit（gcluster job submit 命令）时，请使用 --pathways-proxy-args 设置 XLA 标志：
--pathways-proxy-args="--xla_flag_1=value1 --xla_flag_2=value2"
HLO 转储
如需检查提供给 XLA 编译器的 HLO 输入，您可以使用 --pathways-proxy-args 配置 Pathways，将 HLO 转储到指定的 Cloud Storage 位置：
--pathways-proxy-args="--xla_dump_to=gs://BUCKET_NAME/HLO_PREFIX/"