# Tutorial on itwinai TorchTrainer adapted for the distributed GAN using MNIST dataset

**Author(s)**: Henry Mutegeki (CERN), Matteo Bunino (CERN), Jarl Sondre Sæther (CERN)

The code is adapted from
[this example](https://github.com/pytorch/examples/blob/main/mnist/main.py)
and a simple non-distributed GAN model can be found in a file named `simpleGAN.py`
that serves as a baseline GAN example but focus is mainly on the `train.py`
file for the distributed GAN use case.

## Setup

First, from the root of this repository, build the environment containing
pytorch and deepspeed.
Refer to the [itwinai installation steps](https://itwinai.readthedocs.io/latest/getting-started/getting_started_with_itwinai.html#install-itwinai).

Then navigate to the project working directory

```bash
cd tutorials/distributed-ml/torch-tutorial-GAN
```

## Distributed training on a single node (interactive)

If you want to use SLURM in interactive mode, do the following:

```bash
# Allocate resources
$ salloc --partition=batch --nodes=1 --account=intertwin  --gres=gpu:4 --time=1:59:00
job ID is XXXX
# Get a shell in the compute node (if using SLURM)
$ srun --jobid XXXX --overlap --pty /bin/bash 
# Now you are inside the compute node

# On JSC, you may need to load some modules
ml --force purge
ml Stages/2024 GCC OpenMPI CUDA/12 MPI-settings/CUDA Python HDF5 PnetCDF libaio mpi4py

# ...before activating the Python environment (adapt this to your env name/path)
source ../../../envAI_hdfml/bin/activate
```

To launch the training with torch DDP use:

```bash
torchrun --standalone --nnodes=1 --nproc-per-node=gpu train.py --strategy ddp

# Optional -- from a SLURM login node:
srun --jobid XXXX --ntasks-per-node=1 torchrun --standalone --nnodes=1 --nproc-per-node=gpu train.py --strategy ddp
```

To launch the training with Microsoft DeepSpeed use:

```bash
deepspeed train.py -s deepspeed --deepspeed

# Optional -- from a SLURM login node:
srun --jobid XXXX --ntasks-per-node=1 deepspeed train.py --strategy deepspeed
```

## Distributed training with SLURM (batch mode)

Before running distributed training with slurm, make sure to create the `logs_slurm` folder in your
current working directory to store slurm logs and outputs.

Each distributed strategy has its own SLURM job script, which
should be used to run it:

If you want to distribute the code in `train.py` with **torch DDP**, run from terminal:
  
```bash
export DIST_MODE="ddp"
export RUN_NAME="ddp-itwinai"
export TRAINING_CMD="train.py --strategy ddp"
export PYTHON_VENV="../../../envAI_hdfml"
sbatch --export=ALL,DIST_MODE="$DIST_MODE",RUN_NAME="$RUN_NAME",TRAINING_CMD="$TRAINING_CMD",PYTHON_VENV="$PYTHON_VENV" \
    --job-name="$RUN_NAME-n$N" \
    --output="logs_slurm/job-$RUN_NAME-n$N.out" \
    --error="logs_slurm/job-$RUN_NAME-n$N.err" \
    slurm.sh
```

If you want to distribute the code in `train.py` with **DeepSpeed**, run from terminal:
  
```bash
export DIST_MODE="deepspeed"
export RUN_NAME="deepspeed-itwinai"
export TRAINING_CMD="train.py --strategy deepspeed"
export PYTHON_VENV="../../../envAI_hdfml"
sbatch --export=ALL,DIST_MODE="$DIST_MODE",RUN_NAME="$RUN_NAME",TRAINING_CMD="$TRAINING_CMD",PYTHON_VENV="$PYTHON_VENV" \
    --job-name="$RUN_NAME-n$N" \
    --output="logs_slurm/job-$RUN_NAME-n$N.out" \
    --error="logs_slurm/job-$RUN_NAME-n$N.err" \
    slurm.sh
```

## Analyze the logs

Analyze the logs with MLFlow:

```bash
itwinai mlflow-ui --path mllogs/mlflow
```

## Distributed GAN Documentation

This Guide provides a detailed explanation of how a simple Generative Adversarial Network (GAN) has been adapted to
operate within a distributed environment using the `GANTrainer`. This adaptation enables more efficient training on
larger datasets by leveraging distributed computing resources.

## Overview

A Generative Adversarial Network consists of two key components:

- **Generator (G)**: Generates new data instances.
- **Discriminator (D)**: Evaluates them for authenticity, aiming to distinguish real instances from the fake ones
generated by the Generator.

The training process involves iterative adjustments where the Generator tries to produce data indistinguishable from
actual data, and the Discriminator improves its ability to detect fakes.

## Steps to make a distributed GAN model

The code for all steps can be seen in the attached python file.

### Step 1: Define Model Architecture

Both the Generator and Discriminator are defined using PyTorch's `nn.Module`. The specific architecture for both
includes convolutional layers that are well-suited for processing image data.

### Step 2: Implement Distributed Training

The `GANTrainer` class extends the custom itwinia `TorchTrainer` class and handles the initialization of models,
optimizers, and the distributed training strategy for the GAN. The snippet below shows how the GANTrainer is extending
the TorchTrainer class and initializing the parameters.
This is essentially done to handle the scenario for the GAN which comprises of two Neural Network models which is not
handled by the TorchTrainer that expects and handles one model. We also create custom optimizers for the Optimizer
and Discriminator GAN models:

### Step 3: Training and Validation Logic

The training alternates between updating the Discriminator using real and generated images and training the Generator
to fool the Discriminator. Validation evaluates the performance of the Generator in deceiving the Discriminator.

### Step 4: Visualization and Monitoring

Training progress is monitored through visualizations of loss metrics and image samples generated periodically.

## Takeaways

This readme describes steps taken to adapt a GAN for distributed training, aimed at enhancing efficiency and scalability
for training on large-scale datasets. From this use case we learn the following:

- itwinai `TorchTrainer` can easily be adapted to different unique use cases like the GAN that has two models.
- Training models in a distributed environment may require some level of customization in the training architecturing but
it comes with lots of performance improvements.
- Distributed training for GANs requires a large dataset to reduce the chances of overfitting to the smaller data splits
created during the training phase.
- Always ensure that both the models, data and results for a specific process are accessible on the same device during training
and validation in a distributed environment.
