Ajustement des LLM avec NVIDIA DGX Spark et Unsloth

Tutoriel sur la façon d'affiner et de faire de l'apprentissage par renforcement (RL) avec OpenAI gpt-oss sur NVIDIA DGX Spark.

Unsloth permet l'affinage local de LLMs jusqu'à 200 milliards de paramètres sur le NVIDIA DGX™ Spark. Avec 128 Go de mémoire unifiée, vous pouvez entraîner des modèles massifs tels que gpt-oss-120b, et exécuter ou déployer l'inférence directement sur DGX Spark.

Comme montré à OpenAI DevDay, gpt-oss-20b a été entraîné avec RL et Unsloth sur DGX Spark pour gagner automatiquement à 2048. Vous pouvez entraîner en utilisant Unsloth dans un conteneur Docker ou un environnement virtuel sur DGX Spark.

Dans ce tutoriel, nous allons entraîner gpt-oss-20b avec RL en utilisant les notebooks Unsloth après avoir installé Unsloth sur votre DGX Spark. gpt-oss-120b utilisera environ 68 Go de mémoire unifiée.

Après 1 000 étapes et 4 heures d'entraînement par RL, le modèle gpt-oss surpasse largement l'original sur 2048, et un entraînement plus long améliorerait encore les résultats.

⚡ Tutoriel étape par étape

Commencer avec l'image Docker Unsloth pour DGX Spark

Tout d'abord, construisez l'image Docker en utilisant le Dockerfile DGX Spark qui peut être trouvé ici. Vous pouvez aussi exécuter ce qui suit dans un terminal sur le DGX Spark :

sudo apt update && sudo apt install -y wget
wget -O Dockerfile "https://raw.githubusercontent.com/unslothai/notebooks/main/Dockerfile_DGX_Spark"

Ensuite, construisez l'image Docker d'entraînement en utilisant le Dockerfile sauvegardé :

docker build -f Dockerfile -t unsloth-dgx-spark .

Vous pouvez aussi cliquer pour voir le Dockerfile complet pour DGX Spark

FROM nvcr.io/nvidia/pytorch:25.09-py3

# Définir les variables d'environnement CUDA
ENV CUDA_HOME=/usr/local/cuda-13.0/
ENV CUDA_PATH=$CUDA_HOME
ENV PATH=$CUDA_HOME/bin:$PATH
ENV LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
ENV C_INCLUDE_PATH=$CUDA_HOME/include:$C_INCLUDE_PATH
ENV CPLUS_INCLUDE_PATH=$CUDA_HOME/include:$CPLUS_INCLUDE_PATH

# Installer triton depuis la source pour le support le plus récent de Blackwell
RUN git clone https://github.com/triton-lang/triton.git && \
    cd triton && \
    git checkout c5d671f91d90f40900027382f98b17a3e04045f6 && \
    pip install -r python/requirements.txt && \
    pip install . && \
    cd ..

# Installer xformers depuis la source pour le support de Blackwell
RUN git clone --depth=1 https://github.com/facebookresearch/xformers --recursive && \
    cd xformers && \
    export TORCH_CUDA_ARCH_LIST="12.1" && \
    python setup.py install && \
    cd ..

# Installer unsloth et les autres dépendances
RUN pip install unsloth unsloth_zoo bitsandbytes==0.48.0 transformers==4.56.2 trl==0.22.2

# Lancer le shell
CMD ["/bin/bash"]

Lancer le conteneur

Lancez le conteneur d'entraînement avec accès GPU et montages de volumes :

docker run -it \
    --gpus=all \
    --net=host \
    --ipc=host \
    --ulimit memlock=-1 \
    --ulimit stack=67108864 \
    -v $(pwd):$(pwd) \
    -v $HOME/.cache/huggingface:/root/.cache/huggingface \
    -w $(pwd) \
    unsloth-dgx-spark

Démarrer Jupyter et exécuter les notebooks

Dans le conteneur, démarrez Jupyter et exécutez le notebook requis. Vous pouvez utiliser le notebook Reinforcement Learning gpt-oss 20b pour gagner à 2048 notebook ici. En fait, tous les notebooks Unsloth fonctionnent sur DGX Spark y compris le 120b notebook ! Il suffit de supprimer les cellules d'installation.

Les commandes ci-dessous peuvent également être utilisées pour exécuter le notebook RL. Après le lancement de Jupyter Notebook, ouvrez le “gpt_oss_20B_RL_2048_Game.ipynb”

NOTEBOOK_URL="https://raw.githubusercontent.com/unslothai/notebooks/refs/heads/main/nb/gpt_oss_(20B)_Reinforcement_Learning_2048_Game_DGX_Spark.ipynb"
wget -O "gpt_oss_20B_RL_2048_Game.ipynb" "$NOTEBOOK_URL"

jupyter notebook --ip=0.0.0.0 --port=8888 --no-browser --allow-root

N'oubliez pas qu'Unsloth vous permet aussi de sauvegarder et exécuter vos modèles après l'affinage afin que vous puissiez les déployer localement directement sur votre DGX Spark ensuite.

Un grand merci à Lakshmi Ramesh et Barath Anandan de NVIDIA pour avoir aidé au lancement d'Unsloth sur DGX Spark et pour la construction de l'image Docker.

Utilisation de la mémoire unifiée

L'affinage QLoRA 4 bits de gpt-oss-120b utilisera environ 68 Go de mémoire unifiée. À quoi votre utilisation de la mémoire unifiée devrait ressembler avant (gauche) et après (droite) pendant l'entraînement :