JAX是一个快速发展的 Python 库，用于高性能数值计算和机器学习 (ML) 研究。凭借在药物发现、物理 ML、强化学习和神经图形学中的应用，JAX 在过去几年中得到了令人难以置信的采用。JAX 为开发人员和研究人员提供了许多好处，包括易于使用的 NumPy API 以及与Autograd 的直接集成，以便于区分和优化。JAX 还包括通过几行代码支持跨多节点和多 GPU 系统的分布式处理，并通过NVIDIA GPU 上的XLA优化内核加速性能。
为了帮助更多客户利用 JAX 的强大功能，Oracle 和 NVIDIA 正在合作，以便在 OCI 的多节点集群上轻松设置和使用 JAX。结合 Oracle 云计算平台、NVIDIA GPU 和 RDMA 网络，您可以在云环境中利用 JAX 的可扩展性和性能优势。
在本指南中，我们将介绍如何在 OCI 上设置由NVIDIA A100 Tensor Core GPU提供支持的多节点高性能计算 (HPC) 集群，并内置对 CUDA 和SLURM的支持。我们还将介绍如何将 JAX 安装到该环境中，并开始使用多节点 JAX 代码。
将 GPU + OFED 图像导入 OCI
首先，我们将机器映像导入您的隔间，其中包括对使用 CUDA 的 GPU 计算和使用 OFED 通过以太网进行 RDMA 的支持。
在 Oracle Cloud Console 中，转到菜单，然后在“计算”部分中选择“实例”。
在侧面板中，选择“自定义图像”，然后选择“导入图像”。
选择“从对象存储 URL 导入”和“OCI”。
使用以下 URL 导入图像：
https://objectstorage.ap-osaka-1.oraclecloud.com/p/-O5UHXYmvBSvGcNBhMDH6e263KXVRPit-_wRY6D9275faFJiS2_IiovreESkWtPI/n/hpc_limited_availability/b/Obj_archive/o/OracleLinux-7-UEK-OFED- 5.4-3.1.0.0-GPU-510-2022.05.07-0
该图像可根据当前 OCI 订阅请求访问。
在 OCI 中使用 NVIDIA A100 GPU 创建多节点集群
- 登录到 Oracle 云控制台。
- 在菜单的市场部分，选择“所有应用程序”。
- 搜索并选择“HPC 集群”。
- 选择最新版本和您的 OCI 隔间。然后单击启动堆栈。
（可选）配置您的集群名称、描述和标签。然后，单击“下一步”。
对于您的集群配置，请使用以下设置： - 上传您的 SSH 公钥。如果您没有 SSH 密钥，可以使用 ssh-keygen 并按照提示生成公钥或私钥对。有关更多信息，例如生成密钥或 Windows，请参阅生成 SSH 密钥。
- 对于头节点，选择一个可用性域，并将形状保留为默认值。将磁盘大小增加到 500 GB。
- 对于计算节点，选择 AD 可用性，选择“BM.GPU4.8”形状（8 个 NVIDIA A100 40GB Tensor Core GPU），取消选中超线程，将启动磁盘增加到 500 GB，选择 4 作为“初始集群大小”四个节点，取消选中“使用市场图像”，然后选择您之前导入的图像。
将其他选项保留为默认值，然后单击下一步。在审核页面上，勾选“Run apply”，然后点击Create。
在菜单中，转到“资源管理器”，然后转到“堆栈”。
选择“堆栈详细信息”，然后选择“作业详细信息”。
等待应用作业完成并成功。此过程可能需要 15-45 分钟。
SSH 进入头节点并安装 Python 和 JAX
在菜单中，转到计算，然后转到实例。
找到头节点的公共 IP 地址。
通过以下命令使用 IP 地址通过 SSH 连接到根节点：ssh opc@[IP_ADDRESS]
通过运行以下命令安装 Miniconda 和 Python：
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh && bash ~/miniconda.sh -b -p ${HOME}/miniconda && ~/miniconda/bin/conda
初始化 bash
注销并重新登录。然后使用以下命令安装
Jax：pip3 install — upgrade “jax[cuda]” -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
使用 SLURM 运行您的第一个 JAX 代码
OCI HPC 集群预装了 SLURM，这简化了跨多个节点运行多进程作业。要使用 SLURM 运行 JAX，请创建一个包含以下内容的文件 ~/run.py：
import jax
import jax.numpy as jnp
jax.distributed.initialize()
print(f”Total devices: {jax.device_count()}, “
f”Devices per task: {jax.local_device_count()}”)
x = jnp.ones(jax.local_device_count())
# Computes a reduction (sum) across all devices of x
# and broadcast the result, in y, to all devices.
# If x=[1] on all devices and we have 32 devices,
# the result is y=[32] on all devices.
y = jax.pmap(lambda x: jax.lax.psum(x, “i”), axis_name=”i”)(x)
print(y)
From the login node, run the following command:
(base) [opc@helpful-python-bastion ~]$ srun -N 4 -n 32 — tasks-per-node=8 — gpus-per-node=8 bash -c “~/miniconda/bin/python3 ~/run.py”
You get the following output:
Total devices: 32, Devices per task: 1
[32.]
⋮
Total devices: 32, Devices per task: 1
[32.]
恭喜！您已经成功地在四个节点上运行了一个 JAX 作业，节点上有 32 个 NVIDIA A100 GPU 和总共 32 个进程。有了本指南，您就可以开始多节点 JAX 之旅了。有关 JAX 多主机处理如何工作的更多信息，请访问JAX 文档。有关 JAX GPU 工作负载的更多信息和最新更新，请注册NVIDIA 的 JAX 抢先体验计划，您可以在其中抢先了解 JAX 如何在 GPU 上更快、更大规模地运行。
总结
NVIDIA A100 80GB GPU 的加入补充了 OCI 上可用的 NVIDIA GPU 的现有阵容，包括刚刚发布的 A10 GPU，为世界各地的初创公司、企业和政府在 OCI 上开启了加速计算的新时代。并非所有工作负载都相同，有些工作负载可能需要定制才能在最新一代 GPU 硬件上以最佳方式运行。
