如何在Amazon EC2 G5g实例上配置PyTorch 2.0以启用CUDA检测

核心问题原因 #

G5g实例采用ARM（aarch64）架构，PyTorch的CUDA预编译包针对x86和ARM架构分开发布，你使用的默认cu118索引源仅包含x86架构包，因此pip自动拉取了CPU版本的ARM包。

解决步骤 #

1. 确认系统与硬件环境 #

先验证当前系统架构：

uname -m

输出应为aarch64。同时用nvidia-smi确认GPU驱动和CUDA版本匹配（确保是11.8）。

2. 安装适配ARM架构的CUDA工具链 #

如果是从空白AMI开始，安装aarch64版本的CUDA 11.8：

• 对于Ubuntu 20.04/22.04：

  curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
  curl -s -L https://nvidia.github.io/libnvidia-container/ubuntu20.04/nvidia-container-toolkit.list | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
  sudo apt-get update && sudo apt-get install -y cuda-toolkit-11-8
  

• 对于Amazon Linux 2：

  sudo amazon-linux-extras install nvidia-container-toolkit
  sudo yum install -y cuda-toolkit-11-8
  

安装完成后配置环境变量：

echo 'export PATH=/usr/local/cuda-11.8/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

3. 安装ARM架构的PyTorch 2.0 CUDA版本 #

使用适配aarch64+CUDA的安装命令：

方式一：使用PyTorch官方aarch64 CUDA nightly包

pip3 install torch --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu118 --platform linux_aarch64

方式二：安装AWS优化的PyTorch 2.0版本

AWS针对Graviton编译了优化版PyTorch，可通过以下命令安装：

pip3 install torch==2.0.0 --extra-index-url https://aws-pytorch-binaries.s3.amazonaws.com/torch/2.0.0/cu118

4. 验证安装结果 #

运行以下Python代码确认CUDA可用：

import torch
print(torch.__version__)
print(torch.cuda.is_available())

输出应为True。

额外注意事项 #

• 避免使用x86架构的AMI或安装包，所有操作必须基于aarch64版本的系统镜像。
• 如果使用AWS Deep Learning AMI，选择Graviton版本的镜像，其中已预装适配的CUDA和PyTorch环境，可跳过手动安装步骤。

内容的提问来源于stack exchange，提问作者reed53