Uso de la Nvidia DGX

Para hacer uso de la Nvidia DGX no basta con tener un usuario en el HPC, es necesario solicitar acceso a través del ​CAU. Para saber si un usuario está autorizado a usar la DGX basta con ejecutar el comando sinfo desde el login server. En caso de estarlo, le aparecerá la partición gpu:

PARTITION AVAIL  TIMELIMIT  NODES  STATE NODELIST
normal*      up   infinite      4    mix x440-[01,15,17-18]
normal*      up   infinite      6  alloc x440-[09-14]
normal*      up   infinite     14   idle x440-[02-08,16,19-24]
fast         up   infinite      4   idle x440-[21-24]
fat          up   infinite      1   idle x808
gpu          up   infinite      1   idle averroes-dgx
gpu-bash     up   23:59:00      1   idle averroes-dgx

Aquellos usuarios con acceso a la DGX tendrán en su home del login server un directorio especial llamado dgx cuyo contenido estará almacenado localmente en los discos de la DGX. Solo el contenido de este directorio estará disponible para los trabajos que se manden a la DGX. El motivo de esto es aprovechar los discos de la DGX ya que tienen una mayor velocidad de lectura y escritura que los empleados en Lustre, los cuales pueden suponer un cuello de botellas en tareas de IA como el entrenamiento de redes neuronales.

A continuación se muestra un script sencillo de Slurm que lanza un trabajo en la DGX:

#!/bin/bash
#SBATCH --job-name=dgx_job
#SBATCH --partition=gpu
#SBATCH --chdir=/home/rbarbudo
#SBATCH --gres=gpu:2              # Número de GPUs
#SBATCH --cpus-per-task=4         # Cores de CPU
#SBATCH --mem=16G                 # Memoria RAM
#SBATCH --output=%j.out

. miniconda3/bin/activate
conda activate vllm
python check_gpus.py

Como se puede apreciar es muy similar a los que se utilizan para lanzar trabajos en el clúster y en el fat node, aunque hay una serie de puntos a considerar:

• #SBATCH --partition=gpu: Indica que lo vamos a mandar a la cola en la que los nodos tienen GPUs, en nuestro caso a la DGX
• #SBATCH --gres=gpu:2: Indica el número de GPUs que se reservan. Los jobs que lancemos no podrán acceder a más GPUs de las que se hayan reservado
• #SBATCH --chdir=/home/rbarbudo: Es el directorio en el que vamos a trabajar dentro de la DGX. Hay que tener en cuenta que la ruta /home/<usuario>/dgx que se ve dentro del login server, en la DGX, se corresponde con /home/<usuario>. Por ejemplo, dentro del login server, el fichero check_gpus.py se encuentra dentro de /home/rbarbudo/dgx, aunque dentro de la DGX se encuentra en /home/rbarbudo. Se debe tener en cuenta que cuando mandamos un trabajo a la DGX está solo tiene acceso a los ficheros que están en sus discos.

El script de python que se ejecuta en el shell script es el siguiente:

import torch
import os

def main():
    print("--- Verificación de Recursos en DGX ---")
    
    # Información de Slurm
    job_id = os.environ.get('SLURM_JOB_ID', 'N/A')
    print(f"ID del Trabajo: {job_id}")

    # Verificación de PyTorch y CUDA
    if torch.cuda.is_available():
        device_count = torch.cuda.device_count()
        print(f"GPUs reservadas por Slurm y vistas por PyTorch: {device_count}")
        
        for i in range(device_count):
            props = torch.cuda.get_device_properties(i)
            print(f"  -> GPU {i}: {props.name} ({props.total_memory / 1024**2:.0f} MB)")
    else:
        print("¡ERROR!: PyTorch no detecta GPUs. Verifica la reserva '--gres=gpu:X'.")

if __name__ == "__main__":
    main()

Como se puede apreciar, estamos utilizando la librería torch, muy común en el contexto de IA, para imprimir el número de GPUs accesibles. Si ejecutásemos el script reservando dos GPUs tendríamos la siguiente salida:

--- Verificación de Recursos en DGX ---
ID del Trabajo: 832796
GPUs reservadas por Slurm y vistas por PyTorch: 2
  -> GPU 0: Tesla V100-SXM3-32GB (32501 MB)
  -> GPU 1: Tesla V100-SXM3-32GB (32501 MB)

Abrir una shell en la DGX

Por norma general, toda la carga de datos y la preparación del entorno debe hacerse desde el login node copiando los ficheros necesarios al directorio dgx. No obstante, pueden darse casos en los que sea necesario abrir una shell para, por ejemplo, compilar software en la DGX o instalar un entorno de conda. Para estas situaciones se puede utilizar el comando srun:

srun -p gpu-bash -w averroes-dgx --mem=2GB --qos=gpu-shell --cpus-per-task=1 --pty bash

Con ello abriremos una shell que tendrá asignada un core y 2 GB de RAM. Nótese como en este caso no reservamos ninguna GPU ya que no vamos a realizar ninguna operación que lo requiera. De hecho, si ahora ejecutásemos el script anterior, nos indicaría que no hay ninguna GPU:

rbarbudo@averroes-dgx:~$ . miniconda3/bin/activate 
(base) rbarbudo@averroes-dgx:~$ conda activate vllm
(vllm) rbarbudo@averroes-dgx:~$ python check_gpus.py 
--- Verificación de Recursos en DGX ---
ID del Trabajo: 832837
¡ERROR!: PyTorch no detecta GPUs. Verifica la reserva '--gres=gpu:X'.