Skip to content

Tutorials

Run a containerized PyTorch job

Goal: run a GPU workload from an NGC container image as a normal Slurm batch job. The pattern is pull once, reuse everywhere: convert the image to a SIF file on the shared /home, then point every job at it.

1. Pull the image to a SIF (once)

Run the pull inside a job so the conversion happens on a worker. NGC PyTorch images are big: give the job memory and CPUs (the SIF compression is parallel) and point the Apptainer cache at node-local /local and its temp directory at the per-job /tmp, instead of the slower shared /home. The cache on /local survives the job, so a repeat pull on the same node skips the download. Pick a current NGC PyTorch tag; Blackwell GPUs need a recent build.

srun --mem=64G --cpus-per-task=16 --time=01:00:00 bash -c '
  export APPTAINER_CACHEDIR=/local/apptainer-cache APPTAINER_TMPDIR=/tmp/apptainer-tmp
  mkdir -p "$APPTAINER_CACHEDIR" "$APPTAINER_TMPDIR"
  apptainer pull /home/'"$USER"'/pytorch.sif docker://nvcr.io/nvidia/pytorch:25.06-py3
'

The pull takes about 10 minutes and the resulting SIF is around 13 GB. Because /home is shared and bind-mounted into the jail, the same SIF is now usable from any worker, so this is a one-time step.

2. Submit the job

Write the test script and the batch file, then submit. The batch file references absolute paths on /home (they are filled in when you create the file, so the job itself does not depend on any runtime environment):

cat > /home/$USER/torch-test.py <<'EOF'
import torch
n = torch.cuda.device_count()
print(f"visible GPUs: {n}")
for i in range(n):
    print(" ", torch.cuda.get_device_name(i))
x = torch.randn(8192, 8192, device="cuda", dtype=torch.bfloat16)
y = x @ x
torch.cuda.synchronize()
print("matmul OK:", y.shape)
EOF

cat > /home/$USER/torch-test.sbatch <<EOF
#!/bin/bash
#SBATCH --job-name=torch-test
#SBATCH --nodes=1
#SBATCH --gres=gpu:8
#SBATCH --mem=128G
#SBATCH --time=00:10:00

apptainer exec --nv /home/$USER/pytorch.sif python3 /home/$USER/torch-test.py
EOF

sbatch /home/$USER/torch-test.sbatch

3. Check the result

The job takes under a minute. Watch it with squeue --me, then read its output:

$ cat slurm-<jobid>.out
visible GPUs: 8
  NVIDIA B300 SXM6 AC
  NVIDIA B300 SXM6 AC
  ...
matmul OK: torch.Size([8192, 8192])

Key points:

  • --nv exposes the host NVIDIA driver inside the container; --gres=gpu:8 decides how many GPUs the step (and therefore the container) sees.
  • Request memory (--mem) for both the pull and the job; image conversion and PyTorch both need it.
  • For interactive experimentation, the same works under srun --pty:

    srun --gres=gpu:1 --mem=64G --time=00:30:00 --pty \
      apptainer exec --nv /home/$USER/pytorch.sif python3
    

More container details (GPU passthrough, what is not available) are on the Containers page. For multi-user setups with per-user limits, see the workshop use case.