PyTorch sur ROMEO ARM
Installer PyTorch avec CUDA sur ARM.
Cette documentation décrit l'installation de PyTorch sur les nœuds ARM GPU de Romeo.
Les commandes sont données pour l'environnement Romeo observé lors de la rédaction. Les versions disponibles peuvent évoluer.
1. Obtenir une session interactive sur un nœud ARM GPU
Le nœud de connexion de Romeo utilise l'architecture x86_64, tandis que les nœuds ciblés utilisent l'architecture ARM 64 bits (aarch64).
Il faut donc obtenir une allocation interactive sur un nœud armgpu avant de créer l'environnement Python, d'installer les frameworks ou de compiler.
Consulter la documentation Romeo : Soumission interactive.
1.1. Demander une allocation interactive
Depuis le nœud de connexion :
salloc --time=04:00:00 \
--account=MON_PROJET \
--constraint=armgpu \
--mem=16G \
--gpus=1
Remplacer MON_PROJET par l'identifiant du projet Romeo. Adapter la durée, la mémoire et le nombre de GPU aux besoins et aux quotas disponibles.
Attendre que Slurm accorde l'allocation.
1.2. Ouvrir un shell sur le nœud de calcul
Une fois l'allocation accordée :
srun --pty bash
Vérifier l'architecture :
hostname
uname -m
L'architecture attendue est aarch64.
1.3. Charger l'environnement Spack ARM GPU
Sur le nœud de calcul :
romeo_load_armgpu_env
Ne pas utiliser romeo_load_x64cpu_env, qui correspond à l'environnement CPU x86_64.
Toutes les commandes suivantes doivent être exécutées sur le nœud armgpu.
2. Créer un environnement virtuel Python
Suivre la documentation Romeo : Utiliser Python, via les environnements virtuels Python.
2.1. Choisir une version de Python disponible
Lister les installations Python :
spack find --long python
Exemple de résultat observé sur Romeo :
-- linux-rhel9-neoverse_v2 / no compilers -----------------------
raoq5bs python@3.8.19
74xrf37 python@3.9.19
b6k3ajo python@3.10.14
oxq4fb7 python@3.11.9
a6m2qcc python@3.11.9
zrtbeqk python@3.12.5
rom3mjf python@3.13.0
gzl2pkh python@3.13.0
Pour cet exemple, charger Python 3.13.0 avec le hash gzl2pkh :
spack load py-pip ^python@3.13.0/gzl2pkh
Le hash correspond à une installation présente sur Romeo. Il peut différer selon l'environnement.
Vérifier :
which python
python --version
uname -m
L'interpréteur doit être celui de l'environnement ARM et l'architecture doit être aarch64.
2.2. Créer l'environnement virtuel
mkdir -p ~/mes_environnements_python
python -m venv ~/mes_environnements_python/deep-learning
source ~/mes_environnements_python/deep-learning/bin/activate
Mettre à jour pip et les outils de packaging depuis PyPI :
python -m pip install --upgrade pip setuptools wheel \
--index-url https://pypi.org/simple
Vérifier l'environnement :
which python
python -m pip --version
Un environnement virtuel créé avec un interpréteur x86_64 n'est pas directement utilisable sur aarch64. Il faut le créer depuis l'environnement Python chargé sur le nœud ARM.
3. Vérifier le GPU NVIDIA
Avant d'installer PyTorch, vérifier que le GPU est accessible :
nvidia-smi
Pour afficher uniquement la version CUDA annoncée par le pilote :
nvidia-smi -a | grep CUDA
Exemple observé sur Romeo :
CUDA Version : 13.1
Cette valeur indique la version CUDA maximale prise en charge par le pilote selon nvidia-smi. Elle ne signifie pas que le toolkit CUDA correspondant est chargé.
Pour installer et utiliser le wheel PyTorch précompilé, il n'est pas nécessaire de charger cuda ou cudnn via Spack. Le wheel PyTorch CUDA fournit ses propres bibliothèques d'exécution selon son empaquetage. Le pilote NVIDIA du système reste nécessaire.
4. Installer PyTorch avec CUDA 13.0
4.1. Choisir le wheel ARM
PyTorch publie des wheels précompilés pour certaines combinaisons d'architecture, de Python et de CUDA.
La version CUDA 13.0 a été testée avec succès sur un nœud ARM GPU de Romeo. Le wheel utilisé est compatible avec Python 3.13 et l'architecture aarch64.
Documentation officielle :
4.2. Installer PyTorch
Dans l'environnement virtuel activé :
python -m pip install torch torchvision torchaudio \
--index-url https://download.pytorch.org/whl/cu130 \
--extra-index-url https://pypi.org/simple
L'index cu130 fournit les wheels PyTorch associés à CUDA 13.0. L'index PyPI supplémentaire permet à pip de récupérer les dépendances Python générales.
En cas d'erreur de résolution des dépendances
Sur certaines versions de pip, une installation effectuée uniquement avec l'index PyTorch peut échouer lors de la résolution de dépendances telles que flit_core ou typing-extensions.
L'erreur peut notamment ressembler à :
ERROR: Could not find a version that satisfies the requirement
flit_core<4,>=3.11 (from versions: none)
ERROR: No matching distribution found for flit_core<4,>=3.11
Dans ce cas, mettre d'abord pip à jour depuis PyPI :
python -m pip install --upgrade pip \
--index-url https://pypi.org/simple
Puis relancer l'installation en conservant les deux index :
python -m pip install torch torchvision torchaudio \
--index-url https://download.pytorch.org/whl/cu130 \
--extra-index-url https://pypi.org/simple
Cette méthode a été testée avec succès sur Romeo.
4.3. Vérifier l'installation et le calcul GPU
Exécuter :
python - <<'PY'
import platform
import torch
print("Architecture :", platform.machine())
print("PyTorch :", torch.__version__)
print("CUDA compilée :", torch.version.cuda)
print("GPU disponible :", torch.cuda.is_available())
if torch.cuda.is_available():
print("GPU :", torch.cuda.get_device_name(0))
x = torch.randn(1024, 1024, device="cuda")
print((x @ x).device)
PY
Exemple de résultat obtenu sur un nœud Romeo ARM GPU :
PyTorch : 2.14.1+cu130
CUDA compilée : 13.0
GPU disponible : True
GPU : NVIDIA GH200 120GB
cuda:0
Ces résultats confirment que PyTorch est installé et qu'un calcul matriciel s'exécute sur le GPU.