在 GPU 上執行 SPMD

PyTorch/XLA 支援在 NVIDIA GPU 上執行 SPMD (單節點或多節點)。訓練/推論腳本與用於 TPU 的腳本相同，例如這個 ResNet 腳本。若要使用 SPMD 執行腳本，我們利用 torchrun

PJRT_DEVICE=CUDA \
torchrun \
--nnodes=${NUM_GPU_MACHINES} \
--node_rank=${RANK_OF_CURRENT_MACHINE} \
--nproc_per_node=1 \
--rdzv_endpoint="<MACHINE_0_IP_ADDRESS>:<PORT>" \
training_or_inference_script_using_spmd.py

• --nnodes: 要使用的 GPU 機器數量。

• --node_rank: 目前 GPU 機器的索引。值可以是 0、1、…、${NUMBER_GPU_VM}-1。

• --nproc_per_node: 由於 SPMD 要求，此值必須為 1。

• --rdzv_endpoint: 節點排名 (node_rank)==0 的 GPU 機器的端點，格式為 host:port。host 將會是內部 IP 位址。port 可以是機器上任何可用的連接埠。對於單節點訓練/推論，可以省略此參數。

例如，如果您想在 2 個 GPU 機器上使用 SPMD 訓練 ResNet 模型，您可以在第一部機器上執行以下腳本

XLA_USE_SPMD=1 PJRT_DEVICE=CUDA \
torchrun \
--nnodes=2 \
--node_rank=0 \
--nproc_per_node=1 \
--rdzv_endpoint="<MACHINE_0_INTERNAL_IP_ADDRESS>:12355" \
pytorch/xla/test/spmd/test_train_spmd_imagenet.py --fake_data --batch_size 128

並在第二部機器上執行以下腳本

XLA_USE_SPMD=1 PJRT_DEVICE=CUDA \
torchrun \
--nnodes=2 \
--node_rank=1 \
--nproc_per_node=1 \
--rdzv_endpoint="<MACHINE_0_INTERNAL_IP_ADDRESS>:12355" \
pytorch/xla/test/spmd/test_train_spmd_imagenet.py --fake_data --batch_size 128

如需更多資訊，請參閱 GPU 上 SPMD 支援 RFC。