.. _tutorials-example-03-gpu:

Example 3: PyTorch + MNIST + GPU
================================

This example demonstrates the recommended workflow for running a GPU deep
learning job on the CoSTAR cluster: using **Apptainer** containers.

The following example trains a convolutional neural network on the MNIST
dataset (handwritten digits). It requests a GPU and runs a Python script inside
a public PyTorch container from Docker Hub. This avoids the need to maintain a
Conda installation on the cluster or build your own Apptainer or Docker image
with a Conda environment and PyTorch packages.

1. Connect to the CoSTAR cluster
--------------------------------

First, open your terminal and SSH into the CoSTAR login node:

.. code-block:: bash

   ssh ab1234@costar-login01

Replace ``ab1234`` with your username.

2. Create the submission script
-------------------------------

Download the example submission script, or create a file called
``example-03-gpu.sh`` in your working directory and copy in the script below.

* :download:`example-03-gpu.sh <source_files/example-03-gpu.sh>`

Notice the ``apptainer exec`` command at the bottom of the submission script.
The ``docker://pytorch/pytorch:latest`` image reference tells Apptainer to pull
the public PyTorch container from Docker Hub and run the Python script *inside*
that container.

.. note::

   This example uses the public ``pytorch/pytorch:latest`` Docker image from
   Docker Hub. Apptainer automatically converts the Docker image into an
   Apptainer-compatible image and caches it in the standard Apptainer cache in
   your home directory.

   The first run may take longer while Apptainer pulls and converts the image.
   Later runs are usually faster because the converted image can be reused from
   the cache.

.. dropdown:: example-03-gpu.sh (click to collapse / view)
   :color: primary
   :open:
   :icon: file-code

   .. literalinclude:: source_files/example-03-gpu.sh
      :language: bash
      :linenos:

3. The Python script
--------------------

This is a standard PyTorch training script. Download the example Python script,
or create a file called ``example-03-gpu.py`` in your working directory and
copy in the script below.

* :download:`example-03-gpu.py <source_files/example-03-gpu.py>`

.. dropdown:: example-03-gpu.py (click to collapse / view)
   :color: primary
   :open:
   :icon: file-code

   .. literalinclude:: source_files/example-03-gpu.py
      :language: python
      :linenos:

4. Submit your job
------------------

Submit your script to the Slurm scheduler using the following command:

.. code-block:: bash

   sbatch example-03-gpu.sh

You can check the status of your job by running:

.. code-block:: bash

   squeue --me

5. The output
-------------

Once the job completes, you will see the Slurm output and error files in your
directory:

* **slurm.<node>.<job_id>.out / .err**

  These are the standard Slurm logs. The Python script's training progress,
  including epoch numbers, loss, and accuracy, is written to the ``.out`` file.
  The ``.out`` file should also show that CUDA was detected and name the GPU
  device used by PyTorch.
  If the container failed to pull, or the scheduler encountered an issue, the
  error will be detailed in the ``.err`` file.

.. note::

   If the job is using a GPU and CUDA is available, the output file should
   include lines similar to:

   .. code-block:: text

      CUDA available: True
      CUDA device count: 1
      Using CUDA device: NVIDIA H200
      Train Epoch: 1 [0/60000 (0%)]   Loss: 2.327496
      Train Epoch: 1 [640/60000 (1%)] Loss: 2.328413