---
common_options: &common_options
  data:
    throughput:
      regexp: 'throughput: *(.*?) samples\/sec'
      skip: 2
    loss:
      reduction_type: "final"
      regexp: 'loss: *(\d*\.\d*)'
  output:
    - [samples/sec, "throughput"]
    - [loss, "loss"]
  env:
    PYTORCH_EXE_DIR: "./pt_cache/"

config_options: &config_options
  requirements_path: requirements.txt

validation_options: &validation_options
  description: |
    VIT validation for 4 IPUs.

finetune_options: &finetune_options
  description: |
    VIT finetune training for 16 IPUs.

pretrain_options: &pretrain_options
  description: |
    VIT pretrain training for 16 IPUs.

multihost_options: &multihost_options
  description: |
    ViT pretrain training for 64 IPUs.
    The env variables needed are:
      HOSTS - Comma separated list of hosts, usually something like:
        `10.1.3.101,10.1.3.10x...`
      IPUOF_VIPU_API_HOST - IP of host where VIPU server is running. Usually
        `10.1.3.101`.
      PARTITION - Name of partition to use, can be found via
        `vipu list partitions`. Note that the `--update-partitions` poprun
        flag is disabled here, so ensure the partition is reconfigurable
        or it has been setup properly for this benchmark
      TCP_IF_INCLUDE - Name of subnet, or range of IPs for network interface
        on which all hosts are. Usually `enp65s0f0np0` or `10.1.3.0/24`.

pytorch_vit_finetune_gen_pod16:
  <<: [*common_options, *finetune_options, *config_options]
  description: |
    ViT training for 16 IPUs using generated data
  cmd: >-
    python3 finetune.py
      --config b16_imagenet1k
      --training-steps 10
      --dataset generated
      --checkpoint-output-dir ""
      --executable-cache-dir PYTORCH_EXE_DIR

pytorch_vit_pretrain_gen_pod16:
  <<: [*common_options, *pretrain_options, *config_options]
  cmd: >-
    python3 pretrain.py
      --config b16_in1k_pretrain
      --iterations 10
      --optimizer-state-offchip false
      --byteio true
      --dataset generated
      --epochs 1

pytorch_vit_lamb_pretrain_gen_pod16:
  <<: [*common_options, *pretrain_options, *config_options]
  cmd: >-
    python3 pretrain.py
      --config b16_in1k_pretrain_lamb
      --iterations 10
      --optimizer-state-offchip false
      --byteio true
      --dataset generated
      --epochs 1

pytorch_vit_lamb_pretrain_gen_pod64:
  <<: [*common_options, *multihost_options, *config_options]
  cmd: >-
    poprun
      --vv
      --host $HOSTS
      --ipus-per-replica=4
      --num-ilds=1
      --num-replicas=16
      --num-instances=8
      --reset-partition=no
      --vipu-server-host=$IPUOF_VIPU_API_HOST
      --vipu-server-timeout=3600
      --vipu-partition=$IPUOF_VIPU_API_PARTITION_ID
      --executable-cache-path PYTORCH_EXE_DIR
      --mpi-global-args="
        --mca oob_tcp_if_include $TCP_IF_INCLUDE
        --mca btl_tcp_if_include $TCP_IF_INCLUDE"
      --mpi-local-args="
        -x OPAL_PREFIX
        -x LD_LIBRARY_PATH
        -x PATH
        -x PYTHONPATH
        -x CPATH
        -x IPUOF_VIPU_API_TIMEOUT
        -x POPLAR_LOG_LEVEL
        -x POPLAR_SDK_ENABLED
        -x POPLAR_ENGINE_OPTIONS"
    python3 pretrain.py
      --config b16_in1k_pretrain_lamb
      --dataset generated
      --byteio true
      --iterations 100
      --gradient-accumulation 512
      --epochs 1
      --dataloader-workers 32
      --enable-rts false


pytorch_vit_lamb_pretrain_real_pod64_conv:
  <<: [*common_options, *multihost_options, *config_options]
  cmd: >-
    poprun
      --vv
      --host $HOSTS
      --ipus-per-replica=4
      --num-ilds=1
      --num-replicas=16
      --num-instances=8
      --update-partition=yes
      --remove-partition=yes
      --reset-partition=no
      --vipu-server-host=$IPUOF_VIPU_API_HOST
      --vipu-server-timeout=3600
      --vipu-partition=$IPUOF_VIPU_API_PARTITION_ID
      --vipu-allocation=$VIPU_ALLOCATION_ID
      --executable-cache-path PYTORCH_EXE_DIR
      --mpi-global-args="
        --mca oob_tcp_if_include $TCP_IF_INCLUDE
        --mca btl_tcp_if_include $TCP_IF_INCLUDE"
      --mpi-local-args="
        -x OPAL_PREFIX
        -x LD_LIBRARY_PATH
        -x PATH
        -x PYTHONPATH
        -x CPATH
        -x IPUOF_VIPU_API_TIMEOUT
        -x POPLAR_LOG_LEVEL
        -x POPLAR_SDK_ENABLED
        -x POPLAR_ENGINE_OPTIONS"
    python3 pretrain.py
      --config b16_in1k_pretrain_lamb
      --dataset-path $DATASETS_DIR/imagenet-raw-dataset
      --byteio true
      --gradient-accumulation 512
      --dataloader-workers 32
      --enable-rts false
      --mixup true
      --wandb true
      --wandb-project-name torch-vit
      --wandb-run-name pytorch_vit_lamb_pretrain_real_pod64_conv
      --checkpoint-output-dir output/ckpt-in1k-pretrain


pytorch_vit_pretrain_validation_real_pod4:
  <<: [*validation_options, *config_options]
  cmd: >-
    poprun
      --vv
      --ipus-per-replica=4
      --num-replicas=1
      --num-instances=1
      --update-partition=yes
      --remove-partition=yes
      --reset-partition=no
      --vipu-server-host=$VIPU_CLI_API_HOST
      --vipu-partition=$IPUOF_VIPU_API_PARTITION_ID
      --vipu-allocation=$VIPU_ALLOCATION_ID
      --executable-cache-path PYTORCH_EXE_DIR
    python validation.py
      --config b16_in1k_pretrain_valid
      --dataset-path $DATASETS_DIR/imagenet-raw-dataset
      --pretrained-checkpoint output/ckpt-in1k-pretrain/step_11399
      --wandb true
      --wandb-project-name torch-vit
      --wandb-run-name pytorch_vit_pretrain_validation_real_pod4
      --device-iterations 4
  data:
    accuracy:
      reduction_type: "final"
      regexp: 'Valid Accuracy: *(\d+\.\d+)'
  output:
    - [accuracy, "accuracy"]
  env:
    PYTORCH_EXE_DIR: "./pt_cache/"

pytorch_vit_finetune_real_pod16_conv:
  <<: [*common_options, *finetune_options, *config_options]
  cmd: >-
    poprun
      --vv
      --ipus-per-replica=4
      --num-replicas=4
      --num-instances=1
      --reset-partition=no
      --vipu-server-host=$VIPU_CLI_API_HOST
      --vipu-partition=$IPUOF_VIPU_API_PARTITION_ID
      --vipu-allocation=$VIPU_ALLOCATION_ID
      --executable-cache-path PYTORCH_EXE_DIR
    python3 finetune.py
      --config b16_imagenet1k
      --dataset-path $DATASETS_DIR/imagenet-raw-dataset
      --checkpoint-output-dir ./output/ckpt-in1k-finetune
      --wandb true
      --wandb-project-name torch-vit
      --wandb-run-name pytorch_vit_finetune_real_pod16_conv
      --checkpoint-output-dir output/ckpt-in1k-finetune


pytorch_vit_finetune_validation_real_pod4:
  <<: [*validation_options, *config_options]
  cmd: >-
    poprun
      --vv
      --ipus-per-replica=4
      --num-replicas=1
      --num-instances=1
      --update-partition=yes
      --remove-partition=yes
      --reset-partition=no
      --vipu-server-host=$VIPU_CLI_API_HOST
      --vipu-partition=$IPUOF_VIPU_API_PARTITION_ID
      --vipu-allocation=$VIPU_ALLOCATION_ID
      --executable-cache-path PYTORCH_EXE_DIR
    python validation.py
      --config b16_imagenet1k_valid
      --dataset-path $DATASETS_DIR/imagenet-raw-dataset
      --pretrained-checkpoint output/ckpt-in1k-finetune/step_625
      --wandb true
      --wandb-project-name torch-vit
      --wandb-run-name pytorch_vit_finetune_validation_real_pod4
      --device-iterations 4
  data:
    accuracy:
      reduction_type: "final"
      regexp: 'Valid Accuracy: *(\d+\.\d+)'
  output:
    - [accuracy, "accuracy"]
  env:
    PYTORCH_EXE_DIR: "./pt_cache/"
