From d529eee983fec665e1690b09f4be1b0388379169 Mon Sep 17 00:00:00 2001 From: Anja Virkkunen Date: Wed, 3 Jun 2026 12:51:02 +0300 Subject: [PATCH 1/2] PyTorch container example for ARM GPUs --- triton/usage/gracehopper.rst | 92 ++++++++++++++++++++++++++++++++++++ 1 file changed, 92 insertions(+) diff --git a/triton/usage/gracehopper.rst b/triton/usage/gracehopper.rst index 04c94eb6e..8bb0ade58 100644 --- a/triton/usage/gracehopper.rst +++ b/triton/usage/gracehopper.rst @@ -76,3 +76,95 @@ The hello-world.cu could be e.g.:: return 0; } +Building a simple pytorch environment +------------------------------------- + +Nvidia provides ARM containers for PyTorch, which you can use as a starting point for your own containers. +This example shows how you can extend such a container by installing additional packages from pip. + +A new PyTorch container is built each month, you can browse the selection +in the `Nvidia PyTorch catalog `__. +The following container definition file selects the February 2026 Nvidia PyTorch container with PyTorch +version 2.11 as a starting point: + +.. code-block:: none + + Bootstrap: docker + From: nvcr.io/nvidia/pytorch:26.02-py3 + + %post + pip install transformers==4.57.6 pyyaml==6.0.1 + + %help + An apptainer image based on Nvidia's PyTorch container with ARM CPU architecture. + + The bootstrapped container runs on Ubuntu 24.04, and contains CUDA version 13.1, + OpenMPI 4.1.7, Python 3.12 and PyTorch 2.11. + + This image extends the bootstrapped PyTorch container with transformers package + from HuggingFace. + + PyYAML is a package required by transformers that has already been installed by the + operating system package manager in the container. Pip will try to update PyYAML + (to 6.0.3 at the time this image was created), which will fail the build because + pip cannot change packages installed by the system package manager. + Thus, PyYAML has to be pinned to the version used by the system package manager. + + +You can add other packages you need to the ``pip install`` command in the container definition above. +We also recommend documenting your container in the ``%help`` section +(which packages you have added and why). +Save your container definition to a file (here we will use ``pytorch-transformers-arm.def``) and +track it with version control to back it up for reproducibility. + +Next, you need to build the container image (SIF-file) from the definition file. +Building needs to happen on an ARM-device, which you can achieve for example with +the following sbatch script ``build-pytorch-transformers-arm-container.sh`` like so: + +.. code-block:: slurm + + #!/bin/bash + #SBATCH --job-name=build-arm-container + #SBATCH --partition=gpu-grace-h200-141g + #SBATCH --cpus-per-task=4 + #SBATCH --gpus=1 + #SBATCH --time=01:00:00 + #SBATCH --mem=128G + + # You can replace $WRKDIR with $PWD to create the cache in your current working dir + mkdir -p "$WRKDIR/apptainer_cache" + export APPTAINER_CACHEDIR="$WRKDIR/apptainer_cache" + + apptainer build pytorch-transformers-arm.sif pytorch-transformers-arm.def + +After you have successfully built your container, you can start using it in your scripts. +Here is a simple example of how to run an imaginary Python training script with two arguments +using the container: + +.. code-block:: sh + + #!/bin/bash + #SBATCH --job-name=train-script + #SBATCH --partition=gpu-grace-h200-141g + #SBATCH --cpus-per-task=8 + #SBATCH --gpus=1 + #SBATCH --time=04:00:00 + #SBATCH --mem=256G + + # The --nv argument makes the GPU available within the container + apptainer exec --nv pytorch-transformers-arm.sif \ + python train_script.py \ + --arg1 foo \ + --arg2 bar + +You simply need to prepend calls to your scripts with the apptainer exec command. +For a more comprehensive tutorial on apptainer, please see +`the third lesson `__ of our +`Tuesday Tools & Techniques for HPC (TTT4HPC) course <../../training/scip/ttt4hpc-2024.rst>`__. +Just keep in mind when reading the lesson that it assumes x86 architecture instead of ARM, +so adjust the examples in the tutorial to use ARM. +In other words, be sure to select an ARM container as your starting point, +and run the building script on ARM hardware as shown above. +And if you want or need any help setting up your ARM containers, +you can always join `SciComp garage <../../help/garage.rst>`__ for help. + From 89a32529a37579ae88962f71a5953a6ad186a764 Mon Sep 17 00:00:00 2001 From: Anja Date: Fri, 26 Jun 2026 08:31:04 +0300 Subject: [PATCH 2/2] Fix typo --- triton/usage/gracehopper.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/triton/usage/gracehopper.rst b/triton/usage/gracehopper.rst index 8bb0ade58..65102d3f0 100644 --- a/triton/usage/gracehopper.rst +++ b/triton/usage/gracehopper.rst @@ -76,7 +76,7 @@ The hello-world.cu could be e.g.:: return 0; } -Building a simple pytorch environment +Building a simple PyTorch environment ------------------------------------- Nvidia provides ARM containers for PyTorch, which you can use as a starting point for your own containers.