#----------------------------------------------------------------------------------
defaults: &defaults
  byteio: False
  random_seed: 42
  dataloader_workers: 64
  ipus_per_replica: 4
  synthetic_data: False
  optimizer: SGD
  auto_loss_scaling: False
  weight_decay: 0.0
  recompute_checkpoint_every_layer: True
  attention_probs_dropout_prob: 0.0
  hidden_dropout_prob: 0.1
  layer_norm_eps: 1e-6
  resume_training_from_checkpoint: False
  enable_rts: True
  optimizer_state_offchip: False
  prefetch_depth: 2
  precision: "16.16"
  stochastic_rounding: True
  wandb: False
  representation_size: null
  pretrain: False
  mixup: False
  loss: CELoss
  extra_aug: null
  reduction_type: "mean"
  recompute_mid_layers: [1, 4, 7, 10]

#----------------------------------------------------------------------------------
b16_cifar10: &b16_cifar10
  <<: *defaults

  # Execution
  micro_batch_size: 17
  rebatched_worker_size: 128
  training_steps: 2000
  device_iterations: 1
  replication_factor: 4
  gradient_accumulation: 128
  layers_per_ipu: [3,3,3,3]
  enable_rts: True
  wandb_project_name: "torch-vit-cifar10"

  # Model
  hidden_size: 768
  num_hidden_layers: 12
  num_attention_heads: 12
  matmul_proportion: [0.3, 0.3, 0.3, 0.3]
  mlp_dim: 3072
  patches_size: 16
  num_labels: 10

  # Optimizer
  optimizer: SGD
  warmup_steps: 500
  lr_schedule: cosine
  learning_rate: 0.03
  loss_scaling: 1.0
  weight_decay: 0.0
  momentum : 0.9
  first_order_type: "fp16"

  # Dataset
  dataset: cifar10
  dataset_path: "./data/cifar10"
  pretrained_checkpoint: "google/vit-base-patch16-224-in21k"
  checkpoint_output_dir: "./output/b16_cifar10"
  checkpoint_save_steps: 500

#----------------------------------------------------------------------------------
b16_cifar10_valid: &b16_cifar10_valid
  <<: *defaults

  # Execution
  micro_batch_size: 1
  device_iterations: 1
  replication_factor: 1
  gradient_accumulation: 8
  layers_per_ipu: [3,3,3,3]

  # Model
  hidden_size: 768
  num_hidden_layers: 12
  num_attention_heads: 12
  matmul_proportion: [0.3, 0.3, 0.3, 0.3]
  mlp_dim: 3072
  patches_size: 16
  num_labels: 10

  # Dataset
  dataset: cifar10
  dataset_path: "./data/cifar10"
  pretrained_checkpoint: "./output/b16_cifar10/step_1000"

#----------------------------------------------------------------------------------
b16_imagenet1k: &b16_imagenet1k
  <<: *defaults

  # Execution
  micro_batch_size: 17
  rebatched_worker_size: 256
  training_steps: 625
  device_iterations: 8
  replication_factor: 4
  gradient_accumulation: 30
  layers_per_ipu: [3,3,3,3]
  enable_rts: True
  wandb_project_name: "torch-vit-in1k"

  # Model
  hidden_size: 768
  num_hidden_layers: 12
  num_attention_heads: 12
  matmul_proportion: [0.3, 0.3, 0.3, 0.3]
  mlp_dim: 3072
  patches_size: 16
  num_labels: 1000

  # Optimizer
  optimizer: SGD
  warmup_steps: 100
  lr_schedule: cosine
  learning_rate: 0.08
  loss_scaling: 0.25
  weight_decay: 0.00001
  momentum : 0.9
  first_order_type: "fp16"

  # Dataset
  dataset: imagenet1k
  dataset_path: "./data/imagenet1k/"
  pretrained_checkpoint: "google/vit-base-patch16-224-in21k"
  checkpoint_output_dir: "./output/b16_imagenet1k"
  checkpoint_save_steps: 100

#----------------------------------------------------------------------------------
b16_imagenet1k_ALS:
  <<: *b16_imagenet1k

  # Optimizer
  loss_scaling: 1.0
  auto_loss_scaling: True

#----------------------------------------------------------------------------------
b16_imagenet1k_valid: &b16_imagenet1k_valid
  <<: *defaults

  # Execution
  micro_batch_size: 8
  device_iterations: 1
  replication_factor: 1
  gradient_accumulation: 1
  layers_per_ipu: [3,3,3,3]

  # Model
  hidden_size: 768
  num_hidden_layers: 12
  num_attention_heads: 12
  matmul_proportion: [0.3, 0.3, 0.3, 0.3]
  mlp_dim: 3072
  patches_size: 16
  num_labels: 1000

  # Dataset
  dataset: imagenet1k
  dataset_path: "./data/imagenet1k/"
  pretrained_checkpoint: "./output/b16_imagenet1k/step_625"

#----------------------------------------------------------------------------------
b16_in1k_pretrain: &b16_in1k_pretrain
  <<: *defaults

  # Execution
  micro_batch_size: 8
  # When rebatching with mixup enabled, global batch size needs to be divisible by rebatched_worker_size needs
  rebatched_worker_size: 256
  epochs: 300
  device_iterations: 1
  replication_factor: 4
  gradient_accumulation: 128
  layers_per_ipu: [3,3,3,3]
  enable_rts: True
  wandb: False
  wandb_project_name: "torch-vit-pretrain"
  checkpoint_save_steps: 500
  pretrain: True
  mixup: True
  extra_aug: "imagenet_policy"
  reduction_type: "sum"

  # Model
  hidden_size: 768
  representation_size: 768
  num_hidden_layers: 12
  num_attention_heads: 12
  matmul_proportion: [0.3, 0.3, 0.3, 0.3]
  mlp_dim: 3072
  patches_size: 16
  num_labels: 1000
  attention_probs_dropout_prob: 0.1
  drop_path_rate: 0.0

  # Loss
  alpha: 0.5

  # Optimizer
  optimizer: Adam
  warmup_steps: 10000
  lr_schedule: cosine
  learning_rate: 0.001
  loss_scaling: 128
  weight_decay: 0.004
  momentum : 0.9

  # Dataset
  dataset: imagenet1k
  dataset_path: "./data/imagenet1k/"
  pretrained_checkpoint: ""
  checkpoint_output_dir: "./output/ckpt-in1k-pretrain"

#----------------------------------------------------------------------------------
b16_in1k_pretrain_lamb: &b16_in1k_pretrain_lamb
  <<: *defaults

  # Execution
  micro_batch_size: 8
  rebatched_worker_size: 2048
  training_steps: 11730
  epochs: 600
  device_iterations: 1
  replication_factor: 4
  gradient_accumulation: 2048
  layers_per_ipu: [3,3,3,3]
  enable_rts: True
  wandb: False
  wandb_project_name: "torch-vit-pretrain"
  checkpoint_save_steps: 40
  pretrain: True
  mixup: True
  extra_aug: "imagenet_policy"
  reduction_type: "sum"

  # Model
  hidden_size: 768
  representation_size: 768
  num_hidden_layers: 12
  num_attention_heads: 12
  matmul_proportion: [0.3, 0.3, 0.3, 0.3]
  mlp_dim: 3072
  patches_size: 16
  num_labels: 1000
  attention_probs_dropout_prob: 0.1
  drop_path_rate: 0.0

  # Loss
  alpha: 0.5

  # Optimizer
  optimizer: LAMB
  warmup_steps: 1200
  lr_schedule: cosine
  learning_rate: 0.006
  loss_scaling: 1024
  weight_decay:  0.01
  momentum: 0.9
  adam_betas: [0.9, 0.999]
  adam_eps: 1e-6
  bias_correction: True
  accum_type: "fp32"
  first_order_type: "fp32"
  second_order_type: "fp32"
  max_norm: 65535
  max_norm_bias: 0

  # Dataset
  dataset: imagenet1k
  dataset_path: "./data/imagenet1k/"
  pretrained_checkpoint: ""
  checkpoint_output_dir: "./output/ckpt-in1k-pretrain"
#----------------------------------------------------------------------------------
b16_in1k_pretrain_valid: &b16_in1k_pretrain_valid
  <<: *defaults

  # Execution
  micro_batch_size: 8
  device_iterations: 1
  replication_factor: 1
  gradient_accumulation: 1
  layers_per_ipu: [3,3,3,3]
  pretrain: True

  # Model
  hidden_size: 768
  representation_size: 768
  num_hidden_layers: 12
  num_attention_heads: 12
  matmul_proportion: [0.3, 0.3, 0.3, 0.3]
  mlp_dim: 3072
  patches_size: 16
  num_labels: 1000

  # Dataset
  dataset: imagenet1k
  dataset_path: "./data/imagenet1k/"
  pretrained_checkpoint: "./output/ckpt-in1k-pretrain/step_93599"
