---
# System prepended metadata

title: AI and ML exercises

---

# AI and ML exercises

Throughout these exercises we'll be leveraging the existing ROCm instalation. We can use the existing module to set the environment for it:

```
module purge
module load rocm/6.1.0
``` 

Also, note that these exercises are preapared for MI200 GPU series.

## Setting the virtual environments
These exercises include use cases for PyTorch and TensorFlow using Horovod. Let's prepare the environments to install these frameworks.

We'll be leveraging the system Python instalation, so we'll be creating virtual environments to add the Python packages we need:

```
python3 -m venv --system-site-packages $HOME/venv-pt
python3 -m venv --system-site-packages $HOME/venv-tf
```

## Installing the frameworks
Let's install a PytTorch and Tensrflow suitable to the ROCm version we have available: ROCm 6.1. Two minor versions before or after the current ROCm level should work.

Let's activate our environment for PyTorch.
```
source $HOME/venv-pt/bin/activate 
```
and check the available versions:
```
pip install --index-url https://download.pytorch.org/whl/ torch== |& grep -o '[^ ]*rocm[^ ]*'
pip install --index-url https://download.pytorch.org/whl/ torchvision== |& grep -o '[^ ]*rocm[^ ]*'
pip install --index-url https://download.pytorch.org/whl/ torchaudio== |& grep -o '[^ ]*rocm[^ ]*'
```
It should yield something like:
```
...
2.2.2+rocm5.6
2.2.2+rocm5.7
...
0.17.2+rocm5.6
0.17.2+rocm5.7
...
2.2.2+rocm5.6
2.2.2+rocm5.7
```
Great, we can install the latest versions of each package:
```
pip install --index-url https://download.pytorch.org/whl/ \
  torch==2.2.2+rocm5.7 \
  torchvision==0.17.2+rocm5.7 \
  torchaudio==2.2.2+rocm5.7
```
Quick smoke tests to see if PyTorch can detect all GPUs:
```
python3 -c 'import torch; print("I have this many devices:", torch.cuda.device_count())'
```
I should see `I have this many devices: 8`

Let's install TensorFlow in its respective environment.

```
deactivate
source $HOME/venv-tf/bin/activate 
pip install tensorflow-rocm==
```
We have available TensorFlow 2.14.0 for ROCm 6.0. Let's use that version and do a quick smoke tests to see if TensorFlow detects our GPUs:
```
pip install tensorflow-rocm==2.14.0.600
python3 -c 'from tensorflow.python.client import device_lib ; device_lib.list_local_devices()'
```
It should yield:
```
2024-04-18 11:40:20.107002: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /device:GPU:0 with 63922 MB memory:  -> device: 0, name: AMD Instinct MI210, pci bus id: 0000:63:00.0
2024-04-18 11:40:20.333720: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /device:GPU:1 with 63922 MB memory:  -> device: 1, name: AMD Instinct MI210, pci bus id: 0000:43:00.0
2024-04-18 11:40:20.552465: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /device:GPU:2 with 63922 MB memory:  -> device: 2, name: AMD Instinct MI210, pci bus id: 0000:03:00.0
2024-04-18 11:40:20.778092: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /device:GPU:3 with 63922 MB memory:  -> device: 3, name: AMD Instinct MI210, pci bus id: 0000:26:00.0
2024-04-18 11:40:21.014425: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /device:GPU:4 with 63922 MB memory:  -> device: 4, name: AMD Instinct MI210, pci bus id: 0000:e3:00.0
2024-04-18 11:40:21.235818: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /device:GPU:5 with 63922 MB memory:  -> device: 5, name: AMD Instinct MI210, pci bus id: 0000:c3:00.0
2024-04-18 11:40:21.460124: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /device:GPU:6 with 63922 MB memory:  -> device: 6, name: AMD Instinct MI210, pci bus id: 0000:83:00.0
2024-04-18 11:40:21.683632: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /device:GPU:7 with 63922 MB memory:  -> device: 7, name: AMD Instinct MI210, pci bus id: 0000:a3:00.0
```
We are interested in using Horovod with TensorFlow, so let's install it. Horovod build system was not made ready to ROCm 6.0+, so we need to provide some help to identify the new location for the Cmake files:
```
mkdir -p $HOME/cmake
cat > $HOME/cmake/cmake << EOF
#!/bin/bash -e

if [[ "\$@" == *"--build"* ]] ; then
  $(which cmake) \$@ 
else
  $(which cmake) -DCMAKE_MODULE_PATH=$ROCM_PATH/lib/cmake/hip \$@ 
fi
EOF
chmod +x $HOME/cmake/cmake
```

We can now build using our tuned cmake script:
```
PATH=$HOME/cmake:$PATH \
CPATH=$ROCM_PATH/include/rccl \
HOROVOD_WITHOUT_MXNET=1 \
  HOROVOD_WITHOUT_GLOO=1 \
  HOROVOD_GPU=ROCM \
  HOROVOD_ROCM_HOME=$ROCM_PATH \
  HOROVOD_GPU_OPERATIONS=NCCL \
  HOROVOD_CPU_OPERATIONS=MPI \
  HOROVOD_WITH_MPI=1   \
  HOROVOD_ROCM_PATH=$ROCM_PATH \
  HOROVOD_RCCL_HOME=$ROCM_PATH/include/rccl \
  HOROVOD_RCCL_LIB=$ROCM_PATH/lib \
  HCC_AMDGPU_TARGET=gfx90a,gfx942 \
  HOROVOD_WITH_TENSORFLOW=1 \
  HOROVOD_WITHOUT_PYTORCH=1 \
  pip install --no-cache-dir --force-reinstall --verbose horovod==0.28.1
  ```
Let's define a work directory for us to try some examples.
```
mkdir -p $HOME/ai-with-rocm
```
## PyTorch MNIST example.

MNIST is a quite popular data set for computer vision training. We are fortunate that are many examples on the internet on how to train MNIST dataset and they can usually be run without any changes.

Let's take one of the PyTorch official examples for this - we are training we just two epochs:
```
cd $HOME/ai-with-rocm

deactivate
source $HOME/venv-pt/bin/activate 

curl -LO https://raw.githubusercontent.com/pytorch/examples/main/mnist/main.py
python -u main.py --epochs 2 --batch-size 256


```
We can control which GPU to use with the environmental variable ROCR_VISIBLE_DEVICES and can use the `rocprof` to get more of an idea about the GPU activity:

```
ROCR_VISIBLE_DEVICES=2 \
  rocprof python -u main.py --epochs 1 --batch-size 256
```
The resulting file `results.csv` show the different GPU kernels invoked for this application.

## PyTorch MNIST example - distributed.

We might now be interested in distributing our training accross devices. A way to accomplish this is by taking a distributed data-parallel (DDP) approach where each GPU will train different batch independently and combine the results afterwards. 

There are also several examples on how to do this. We will use as starting example https://raw.githubusercontent.com/kubeflow/examples/master/pytorch_mnist/training/ddp/mnist/mnist_DDP.py.

```
cd $HOME/ai-with-rocm
curl -LO https://raw.githubusercontent.com/kubeflow/examples/master/pytorch_mnist/training/ddp/mnist/mnist_DDP.py
```

Download a modified version of this script and compare the differences:

```
curl -LO https://raw.githubusercontent.com/amd/HPCTrainingExamples/main/MLExamples/mnist_DDP_modified.py
diff mnist_DDP.py mnist_DDP_modified.py
```

There are a couple of differences one controls the batch size another one controls how the distributed run is initialized.
```
dist.init_process_group(backend='nccl', 
                        init_method='env://', 
                        world_size=int(os.environ['WORLD_SIZE']),                             rank=int(os.environ['RANK']))
```

PyTorch provides an object to control the distributed run environment:
```
import torch.distributed as dist
```
Here we are instructing that we want to use RCCL (AMD implementation for NCCL) and also want the tool to leverage the environment to collect more information, with some explicit information about the number of ranks (world size) and rank.

Other relevant bits to enable distributed run are in:
```
def run(modelpath, gpu):
  ...
  model = Net()
  ...
  model = torch.nn.parallel.DistributedDataParallel(model)
  ...
  optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
  ...
```
Here the model is wrapped into the `DistributedDataParallel` object to enable the distributed training.

Now to train the model on multiple ranks we will leverage MPI to start the processes and translate the MPI environment to something that PyTorch distributed package understands:

```
cat > run-me.sh << EOF
#!/bin/bash -e
  export MASTER_ADDR=localhost
  export MASTER_PORT=29500
  export WORLD_SIZE=\$OMPI_COMM_WORLD_SIZE
  export RANK=\$OMPI_COMM_WORLD_RANK
  export ROCR_VISIBLE_DEVICES=\$OMPI_COMM_WORLD_LOCAL_RANK
  
  python -u mnist_DDP_modified.py \
    --gpu --modelpath $HOME/ai-with-rocm/model

EOF
chmod +x run-me.sh

mpirun -np 2 ./run-me.sh
```

Master address and port are defined for the different ranks to communicate between themselves. We then leverage the `OMPI_*` variables to decide ranks and GPUs to be used by each of these ranks.

Another popular way to spin a distributed run is to leverage the `torchrun` utility. However, this requires the application to include logic to decide which GPU to use, instead of relying on `ROCR_VISIBLE_DEVICES`. We can add the following line to our application after `dist.init_process_group()` to accomplish that:

```
torch.cuda.set_device(int(os.environ['RANK']))
```
This will make GPUs to be indexed by the rank.

We can then run our application with `torchrun` as:
```
ROCR_VISIBLE_DEVICES=0,1 \
  torchrun --nnodes 1 --nproc_per_node 2 \
    ./mnist_DDP_modified.py --gpu --modelpath $HOME/ai-with-rocm/model
```

## TensorFlow with Horovod example

Similarly to PyTorch, TensorFlow examples should not need changes due to the GPU architecture.

Let's pick up TensorFlow example from the Horovod project - a syntectic training example already rigged to use Horovod:
```
deactivate
source $HOME/venv-tf/bin/activate 
cd $HOME/ai-with-rocm

curl -LO https://raw.githubusercontent.com/horovod/horovod/master/examples/tensorflow2/tensorflow2_synthetic_benchmark.py

mpirun -np 2 \
  python -u tensorflow2_synthetic_benchmark.py --batch-size 256
```

Horovod makes it rather straingforward to start a distributed learning model as it leverages the already existing MPI environment.

We can inspect the test case and see the relevant bits where Horovod `hvd` is being leveraged, namely the selection of the GPU based on the local rank and wrapping of the local Gradient Tape operator.
```
import tensorflow as tf
import horovod.tensorflow as hvd
...
tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')
...
    with tf.GradientTape() as tape:
...
    tape = hvd.DistributedGradientTape(tape, compression=compression)
```

We can further inspect the GPU activity by leveraging `rocprof` to obtain a trace for one of the ranks. We'll do just 2 iterations so that the result files are not too large. 
```
mpirun -np 2 \
  bash -c 'if [ $OMPI_COMM_WORLD_RANK -eq 1 ] ; then \
             profiler="rocprof --hip-trace" ; \
           fi ; \
           $profiler python -u tensorflow2_synthetic_benchmark.py \
             --batch-size 256 \
             --num-warmup-batches 2 \
             --num-iters 2'
            
```
We will get a `results.json` file with the trace information to load in the Perfetto tool. It is often a good idea to compress the file to make it easier to copy to one's workstation.
```
xz -T8 -9 results.json
```
We can then copy it and visualize in Perfetto. We can detect the kernesl from the different libraries, like MIOpen, rocBLAS, Eigen as well as MLIR JIT kernels.
![image](https://hackmd.io/_uploads/B1SpcGyW0.png)

## Examples with Huggingface transformers

There are several repositories of examples. A popular one is the Huggingface transformers. These examples should just work without any modification specific to AMD GPUs.

We can install the transformaer package from source as:
```
deactivate
source $HOME/venv-tf/bin/activate 

git clone \
https://github.com/huggingface/transformers.git \
  $HOME/ai-with-rocm/transformers
cd $HOME/ai-with-rocm/transformers
pip3 install -e .
```

It is useful to point the implementation to a suitable plce to store and cache information and datasets. That can be done by setting the environment variables: 
```
export HF_HOME=$HOME/ai-with-rocm/hf-home
export HUGGINGFACE_HUB_CACHE=$HOME/ai-with-rocm/hf-cache
mkdir -p $HF_HOME $HUGGINGFACE_HUB_CACHE 
```
We can now try the examples. 

### Image calssification
Let's look at the image classification one. We need first to install the example dependences:
``` 
cd $HOME/ai-with-rocm/transformers/examples/pytorch/image-classification
pip3 install -r requirements.txt
pip3 install scikit-learn
pip3 install -U pillow 
```
We are now ready to run the example. We'll be experimenting with using mixed precision (with BF16 datatypes) or not (the default):
```
for precision in '' '--bf16' ; do
  ROCR_VISIBLE_DEVICES=2 \
    python3 run_image_classification.py \
      --dataset_name beans \
      --label_column_name labels \
      --output_dir $HOME/ai-with-rocm/hf-output  \
      --overwrite_output_dir \
      --remove_unused_columns False \
      --do_train \
      --learning_rate 2e-5 \
      --num_train_epochs 2 \
      --per_device_train_batch_size 8 \
      --torch_compile True \
      --seed 1337 \
      $precision
done
```
Depending on what is bounding the performance it could be good idea to use mixed precision or not - so this is a fair experiment to do. This should yield results like:

```
...
# No mixed precision
***** train metrics *****                                                               
  epoch                    =         2.0
  total_flos               = 149248978GF
  train_loss               =      0.4118
  train_runtime            =  0:01:17.52
  train_samples_per_second =      26.675
  train_steps_per_second   =       3.354
...
# With mixed precision (BF16)
***** train metrics *****
  epoch                    =         2.0
  total_flos               = 149248978GF
  train_loss               =      0.4126
  train_runtime            =  0:01:23.73
  train_samples_per_second =      24.698
  train_steps_per_second   =       3.105
```

### Language modeling
A growing set of applications for distributed learning is language modelling. A typical approach is to leverage an established model and fine tune it for one needs. That's what the `question-answering` example does. We'll start our fine-tuning from the BERT base dataset. The steps are similar to what we did before: first install the requirements:

```
cd $HOME/ai-with-rocm/transformers/examples/pytorch/question-answering
pip3 install -r requirements.txt
```
Then, we are ready to run the fine-tuning, e.g., on 2 GPUs for different training precisions:
```
for precision in '' '--bf16' '--fp16' ; do
ROCR_VISIBLE_DEVICES=0,1 \
  torchrun --nproc_per_node 2  run_qa.py \
    --model_name_or_path bert-base-uncased \
    --dataset_name squad \
    --do_train \
    --per_device_train_batch_size 12 \
    --learning_rate 3e-5 \
    --num_train_epochs 1 \
    --max_seq_length 384 \
    --doc_stride 128 \
    --output_dir $HOME/ai-with-rocm/hf-output2  \
    --torch_compile True \
    --overwrite_output_dir \
    $precision
done
```



```
# FP32
***** train metrics *****
  epoch                    =        1.0
  total_flos               = 16159030GF
  train_loss               =     1.3039
  train_runtime            = 0:10:09.70
  train_samples            =      88524
  train_samples_per_second =    145.192
  train_steps_per_second   =      6.051
 
 # BF16
***** train metrics *****
  epoch                    =        1.0
  total_flos               = 16159030GF
  train_loss               =     1.3013
  train_runtime            = 0:07:34.81
  train_samples            =      88524
  train_samples_per_second =    194.636
  train_steps_per_second   =      8.111
 
 # FP16
 ***** train metrics *****
  epoch                    =        1.0
  total_flos               = 16159030GF
  train_loss               =     1.3064
  train_runtime            = 0:06:35.75
  train_samples            =      88524
  train_samples_per_second =    223.686
  train_steps_per_second   =      9.322
 ```
 
## Note about MI300

For MI300 series, the publically available wheels files might not contain the support for it yet. Though this is gahnging eminently, you might have to, e.g. leverage the files in https://repo.radeon.com/rocm/manylinux/. E.g. to install PyTorch and run the MNIST example above we could do:

```
deactivate

python3 -m venv --system-site-packages $HOME/venv-pt-mi300
source $HOME/venv-pt-mi300/bin/activate 

pip install https://repo.radeon.com/rocm/manylinux/rocm-rel-6.0.2/torch-2.1.2%2Brocm6.0-cp310-cp310-linux_x86_64.whl
pip install https://repo.radeon.com/rocm/manylinux/rocm-rel-6.0.2/torchvision-0.16.1%2Brocm6.0-cp310-cp310-linux_x86_64.whl
pip install transformers

cd $HOME/ai-with-rocm
curl -LO https://raw.githubusercontent.com/pytorch/examples/main/mnist/main.py
python -u main.py --epochs 2 --batch-size 256
```