|
| 1 | +# TensorFlow |
| 2 | + |
| 3 | +:::info |
| 4 | +This was adapted from [Princeton University Multi-GPU Training with PyTorch](https://github.com/PrincetonUniversity/multi_gpu_training) |
| 5 | +::: |
| 6 | + |
| 7 | +The starting point for [multi-GPU training with Keras](https://www.tensorflow.org/tutorials/distribute/keras) is `tf.distribute.MirroredStrategy`. In this approach, the model is copied to `N` GPUs and gradients are synced as we saw previously. Be sure to use [`tf.data`](https://www.tensorflow.org/api_docs/python/tf/data) to handle data loading as is done in the example on this page and is explained graphically [here](https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/data_performance.ipynb#scrollTo=i3NtGI3r-jLp). |
| 8 | + |
| 9 | +## Single-Node, Synchronous, Multi-GPU Training |
| 10 | + |
| 11 | +Here were train the ResNet-50 model on the Cassava dataset (see [video](https://www.youtube.com/watch?v=xzSCvXDcX68) on TensorFlow YouTube channel). Here is another example [video](https://www.youtube.com/watch?v=HCLmM1PyDIs) using the "cats vs. dog" dataset. |
| 12 | + |
| 13 | +### Step 1: Create a TensorFlow Overlay File |
| 14 | + |
| 15 | +```bash |
| 16 | +# create a working directory |
| 17 | +[NetID@log-1 ~]$ mkdir /scratch/<NetID>/tensorflow-example |
| 18 | +[NetID@log-1 ~]$ cd /scratch/<NetID>/tensorflow-example |
| 19 | + |
| 20 | +# copy over an overlay file with sufficient resources and unzip it |
| 21 | +[NetID@cm001 tensorflow-example]$ cp -rp /scratch/work/public/overlay-fs-ext3/overlay-15GB-500K.ext3.gz . |
| 22 | +[NetID@cm001 tensorflow-example]$ gunzip overlay-15GB-500K.ext3.gz |
| 23 | + |
| 24 | +# start the singularity environment |
| 25 | +[NetID@cm001 tensorflow-example]$ singularity exec --overlay overlay-15GB-500K.ext3:rw /scratch/work/public/singularity/cuda12.1.1-cudnn8.9.0-devel-ubuntu22.04.2.sif /bin/bash |
| 26 | + |
| 27 | +# install miniforge in singularity environment |
| 28 | +Singularity> wget --no-check-certificate https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh |
| 29 | +Singularity> bash Miniforge3-Linux-x86_64.sh -b -p /ext3/miniforge3 |
| 30 | + |
| 31 | +# create an miniconda environment file |
| 32 | +Singularity> touch /ext3/env.sh |
| 33 | +Singularity> nano /ext3/env.sh |
| 34 | +# and add the following content to it: |
| 35 | +``` |
| 36 | + |
| 37 | +```bash |
| 38 | +#!/bin/bash |
| 39 | + |
| 40 | +unset -f which |
| 41 | + |
| 42 | +source /ext3/miniforge3/etc/profile.d/conda.sh |
| 43 | +export PATH=/ext3/miniforge3/bin:$PATH |
| 44 | +export PYTHONPATH=/ext3/miniforge3/bin:$PATH |
| 45 | +``` |
| 46 | + |
| 47 | +```bash |
| 48 | +# activate the new environment and initialize |
| 49 | +Singularity> source /ext3/env.sh |
| 50 | +Singularity> conda config --remove channels defaults |
| 51 | +Singularity> conda clean --all --yes |
| 52 | +Singularity> conda update -n base conda -y |
| 53 | +Singularity> conda install pip -y |
| 54 | +Singularity> conda install ipykernel -y |
| 55 | + |
| 56 | +# check that you're using the miniconda environment |
| 57 | +# you should get the following output |
| 58 | +Singularity> which conda |
| 59 | +/ext3/miniforge3/bin/conda |
| 60 | +Singularity> which python |
| 61 | +/ext3/miniforge3/bin/python |
| 62 | +Singularity> which pip |
| 63 | +/ext3/miniforge3/bin/pip |
| 64 | + |
| 65 | +# install TensorFlow |
| 66 | +Singularity> pip install tensorflow[and-cuda] tensorflow_datasets |
| 67 | +``` |
| 68 | + |
| 69 | +### Step 2: Download the Data |
| 70 | + |
| 71 | +This example using the `cassava` dataset which requires 4 GB of storage space. Be sure to save this in your `/scratch` space and not in `/home`. |
| 72 | + |
| 73 | +Please save the following into a file named `download_data_and_weights.py`: |
| 74 | +```python |
| 75 | +import tensorflow as tf |
| 76 | +import tensorflow_datasets as tfds |
| 77 | + |
| 78 | +# download the data (4 GB) on the login node |
| 79 | +_ = tfds.load(name='cassava', with_info=True, as_supervised=True, data_dir='.') |
| 80 | + |
| 81 | +# download the model weights on the login node |
| 82 | +_ = tf.keras.applications.ResNet50(weights="imagenet", include_top=False) |
| 83 | +``` |
| 84 | + |
| 85 | +Run the command below to download the data (4 GB in size): |
| 86 | + |
| 87 | +```bash |
| 88 | +# switch to a data transfer node |
| 89 | +[NetID@log-1 tensorflow_example]$ ssh gdtn |
| 90 | +[NetID@dtn-1 ~]$ cd /scratch/NetID/tensorflow_example |
| 91 | +[NetID@dtn-1 tensorflow_example]$ singularity exec --nv --overlay /scratch/NetID/pytorch-example/my_pytorch.ext3:ro /scratch/work/public/singularity/cuda12.1.1-cudnn8.9.0-devel-ubuntu22.04.2.sif /bin/bash -c "source /ext3/env.sh; python download_data_and_weights.py" |
| 92 | +``` |
| 93 | + |
| 94 | +### Step 3: Inspect the Script |
| 95 | + |
| 96 | +Below is the contents of `mnist_classify.py`: |
| 97 | + |
| 98 | +```python |
| 99 | +import argparse |
| 100 | +import os |
| 101 | +import tensorflow_datasets as tfds |
| 102 | +import tensorflow as tf |
| 103 | +from time import perf_counter |
| 104 | + |
| 105 | +def preprocess_data(image, label): |
| 106 | + image = tf.image.resize(image, (300, 300)) |
| 107 | + image = tf.cast(image, tf.float32) / 255.0 |
| 108 | + return image, label |
| 109 | + |
| 110 | +def create_dataset(batch_size_per_replica, datasets, strategy): |
| 111 | + batch_size = batch_size_per_replica * strategy.num_replicas_in_sync |
| 112 | + return datasets['train'].map(preprocess_data, num_parallel_calls=tf.data.AUTOTUNE) \ |
| 113 | + .cache() \ |
| 114 | + .shuffle(1000) \ |
| 115 | + .batch(batch_size) \ |
| 116 | + .prefetch(tf.data.AUTOTUNE) |
| 117 | + |
| 118 | +def create_model(num_classes): |
| 119 | + base_model = tf.keras.applications.ResNet50(weights="imagenet", include_top=False) |
| 120 | + x = base_model.output |
| 121 | + x = tf.keras.layers.GlobalAveragePooling2D()(x) |
| 122 | + x = tf.keras.layers.Dense(1016, activation="relu")(x) |
| 123 | + predictions = tf.keras.layers.Dense(num_classes, activation="softmax")(x) |
| 124 | + model = tf.keras.Model(inputs=base_model.input, outputs=predictions) |
| 125 | + return model |
| 126 | + |
| 127 | +def train(epochs, num_classes, train_dataset, strategy): |
| 128 | + with strategy.scope(): |
| 129 | + model = create_model(num_classes) |
| 130 | + model.compile(loss='sparse_categorical_crossentropy', |
| 131 | + optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001), |
| 132 | + metrics=['accuracy']) |
| 133 | + |
| 134 | + start_time = perf_counter() |
| 135 | + model.fit(train_dataset, epochs=epochs) |
| 136 | + print("Training time:", perf_counter() - start_time) |
| 137 | + return None |
| 138 | + |
| 139 | +def print_info(num_replicas_in_sync, batch_size_per_replica, info, num_classes): |
| 140 | + print(f'TF Version: {tf.__version__}') |
| 141 | + print(f'Number of GPUs: {num_replicas_in_sync}') |
| 142 | + print(f'Batch size per GPU: {batch_size_per_replica}') |
| 143 | + print(f'Train records: {info.splits["train"].num_examples}') |
| 144 | + print(f'Test records: {info.splits["test"].num_examples}') |
| 145 | + print(f'Number of classes: {num_classes}') |
| 146 | + return None |
| 147 | + |
| 148 | +if __name__ == '__main__': |
| 149 | + parser = argparse.ArgumentParser(description='Multi-GPU Training Example') |
| 150 | + parser.add_argument('--batch-size-per-replica', type=int, default=32, metavar='N', |
| 151 | + help='input batch size per GPU for training (default: 32)') |
| 152 | + parser.add_argument('--epochs', type=int, default=15, metavar='N', |
| 153 | + help='number of epochs to train (default: 15)') |
| 154 | + args = parser.parse_args() |
| 155 | + |
| 156 | + datasets, info = tfds.load(name='cassava', with_info=True, as_supervised=True, data_dir=".") |
| 157 | + num_classes = info.features["label"].num_classes |
| 158 | + |
| 159 | + strategy = tf.distribute.MirroredStrategy() |
| 160 | + train_dataset = create_dataset(args.batch_size_per_replica, datasets, strategy) |
| 161 | + train(args.epochs, num_classes, train_dataset, strategy) |
| 162 | + |
| 163 | + print_info(strategy.num_replicas_in_sync, args.batch_size_per_replica, info, num_classes) |
| 164 | +``` |
| 165 | + |
| 166 | +### Step 4: Submit the Job |
| 167 | + |
| 168 | +Below is a sample Slurm script: |
| 169 | + |
| 170 | +```bash |
| 171 | +#!/bin/bash |
| 172 | +#SBATCH --job-name=cassava # create a short name for your job |
| 173 | +#SBATCH --nodes=1 # node count |
| 174 | +#SBATCH --ntasks=1 # total number of tasks across all nodes |
| 175 | +#SBATCH --cpus-per-task=16 # cpu-cores per task (>1 if multi-threaded tasks) |
| 176 | +#SBATCH --mem=64G # total memory per node (4G per cpu-core is default) |
| 177 | +#SBATCH --gres=gpu:2 # number of gpus per node |
| 178 | +#SBATCH --time=00:20:00 # total run time limit (HH:MM:SS) |
| 179 | + |
| 180 | +module purge |
| 181 | + |
| 182 | +srun singularity exec --nv \ |
| 183 | + --overlay /scratch/NetID/pytorch_examples_new/tensorflow-example/tensorflow.ext3:ro \ |
| 184 | + /scratch/work/public/singularity/cuda12.1.1-cudnn8.9.0-devel-ubuntu22.04.2.sif\ |
| 185 | + /bin/bash -c "source /ext3/env.sh; python mnist_classify.py --batch-size-per-replica=32 --epochs=15" |
| 186 | +``` |
| 187 | +:::note |
| 188 | +Be sure to change `NetID` in the above script to your NetID. |
| 189 | +::: |
| 190 | + |
| 191 | +Submit the job as follows: |
| 192 | + |
| 193 | +```bash |
| 194 | +[NetID@log-1 tensorflow_example]$ sbatch job.slurm |
| 195 | +``` |
| 196 | + |
| 197 | +### Performance |
| 198 | + |
| 199 | +The training time is shown below for different choices of `cpus-per-task` and the number of GPUs on a test system (your results will vary depending on your system specs): |
| 200 | + |
| 201 | +| nodes | ntasks | cpus-per-task | GPUs | Training Time (s) | Mean GPU Utilization (%) | |
| 202 | +|:-------------:|:-------------:|:------------:|:--------:|:-----------------:|:-------------------------:| |
| 203 | +| 1 | 1 | 2 | 1 | 574 | 85 | |
| 204 | +| 1 | 1 | 4 | 1 | 565 | 83 | |
| 205 | +| 1 | 1 | 8 | 1 | 562 | 89 | |
| 206 | +| 1 | 1 | 16 | 1 | 564 | 90 | |
| 207 | +| 1 | 1 | 4 | 2 | 339 | 76 | |
| 208 | +| 1 | 1 | 8 | 2 | 334 | 81 | |
| 209 | +| 1 | 1 | 16 | 2 | 332 | 74 | |
| 210 | +| 1 | 1 | 4 | 3 | 256 | 68 | |
| 211 | +| 1 | 1 | 8 | 3 | 251 | 73 | |
| 212 | +| 1 | 1 | 16 | 3 | 249 | 66 | |
| 213 | +| 1 | 1 | 4 | 4 | 226 | 59 | |
| 214 | +| 1 | 1 | 8 | 4 | 220 | 58 | |
| 215 | +| 1 | 1 | 16 | 4 | 214 | 65 | |
| 216 | +| 1 | 1 | 32 | 4 | 218 | 63 | |
| 217 | + |
| 218 | +"Training Time" in the table above is the time to run `model.fit(train_dataset, epochs=epochs)`. "Mean GPU utilization" was taken from the output of the `jobstats` command. |
| 219 | + |
| 220 | +The figure below shows the speed-up as a function of the number of GPUs. The dashed line shows the maximum possible speed-up. |
| 221 | + |
| 222 | + |
| 223 | + |
| 224 | +We see that linear scaling is not observed. That is, the training time when using 2 GPUs is not 1/2 of the training time when using one. To improve on this one would profile the script and identify the performance bottleneck. Some of the training images are 500 pixels wide. It could be that the preprocessing step is the slowest. |
| 225 | + |
| 226 | +## Multi-node Training |
| 227 | + |
| 228 | +- [MultiWorkerMirroredStrategy](https://www.tensorflow.org/guide/distributed_training#multiworkermirroredstrategy) for using the GPUs on more than one compute node. |
| 229 | +- [Keras API](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras). |
| 230 | +- [Horovod](https://horovod.ai/) is a distributed deep learning training framework for TensorFlow, Keras, PyTorch, and Apache MXNet. It is based on MPI. |
0 commit comments