launch.py

#!/usr/bin/env python
"""Launch training on AWS with 8 GPUs."""

import argparse
import re

import ncluster
from attrdict import AttrDefault

import util

parser = argparse.ArgumentParser()
parser.add_argument('--name', type=str, default='txl',
                    help="name of the current run, used for machine naming and rundir name")
parser.add_argument('--config', type=str, default='',
                    help='which training config to use')
parser.add_argument('--nospot', action='store_true',
                    help='Use more expensive on-demand instance')

parser.add_argument('--skip_setup', action='store_true',
                    help='Make startup slightly faster by skiping various initialization tasks, like '
                         'tmux/efs setup. Only use on reruns.')

parser.add_argument('--wiki', action='store_true',
                    help='Train on all of wikipedia.')
parser.add_argument('--git', action='store_true',                    help='Train on git dataset.')
parser.add_argument('--bpe', action='store_true',
                    help='Use BPE to reduce vocab instead of adaptive softmax div')

# network settings
parser.add_argument('--num_rings', type=int, default=16)

# config flags for trying configurations outside of existing configs
parser.add_argument('--machines', type=int, default=0,
                    help="how many machines to use")
parser.add_argument('--instance_type', type=str, default='',
                    help="how many machines to use")
parser.add_argument('--checkpoint_each_epoch', type=int, default=0,
                    help='whether to save checkpoint at each epoch')
parser.add_argument('--image_name', type=str, default='',
                    help="use custom AMI ")
parser.add_argument('--conda_env', type=str, default='',
                    help='use custom conda env')
parser.add_argument('--checkpoint', type=str, default='',
                    help='restore from this checkpoint')
parser.add_argument('--optim_state_dict', type=str, default='',
                    help='restore optimizer from this checkpoint')

parser.add_argument('--skip_files', type=float, default=0,
                    help='how many files skip in the first epoch')
parser.add_argument('--valid_custom', type=str, default=None,
                    help='url to custom valid file')

args = parser.parse_args()

# default environment settings, should change rarely since they affect
# all configs
IMAGE_NAME = 'cybertronai01_git04'
# CONDA_ENV = 'pytorch_april_nccl237'
CONDA_ENV = 'pytorch_p36'

# 'base_lr': learning rate for BASE_LR_BATCHSIZE, linear lr scaling will grow this rate proportionally to final global
# batch size
# local_batch_size: per-GPU batch size
BASE_LR_BATCHSIZE = 32

# logs: yaro-1gpu
one_gpu = {
    # 24x smaller batch than ben-big-lr.09, use 5x more agressive learning rate
    'base_lr': 0.000125 * 5 / 3 * 5,
    'local_batch_size': 32,
    'instance_type': 'p3.2xlarge',
    'machines': 1,
}

# logs: yaro-one.08
one_machine = {
    'base_lr': 0.000125 * 5 / 3,  # from ben-big-lr.09
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 96,
    'machines': 1,
}

# logs: yaro-fp16
one_machine_fp16 = {
    'base_lr': 0.000125 * 5 / 3,  # from ben-big-lr.09
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 96,
    'machines': 1,
    'extra_worker_params': {
        'fp16': True,
        'dynamic_loss_scale': True,
    }
}

one_small_machine = {
    'base_lr': 0.001 / 4,  # Divide by 4 to counteract batch adjustment
    'instance_type': 'p3.16xlarge',
    'local_batch_size': 6,
    'machines': 1,
    'large': True,
    'extra_worker_params': {
        'fp16': True,
        'dynamic_loss_scale': True,
    }
}

one_small_machine_wiki = {
    'base_lr': 0.001 / 4,  # Divide by 4 to counteract batch adjustment
    'instance_type': 'p3.16xlarge',
    'local_batch_size': 6,
    'machines': 1,
    'large': True,
    'checkpoint': '/ncluster/runs/txl.09/model-best.pt',  # us-east-1
    'optim_state_dict': '/ncluster/runs/txl.09/optimizer-best.pt',
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e6,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/wikiextracted',
        'dataset': 'wiki',
    }
}

one_tiny_machine_git = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3.2xlarge',
    'local_batch_size': 6,
    'machines': 1,
    'large': True,
    # 'checkpoint': '/ncluster/runs/txl.09/model-best.pt',  # us-east-1
    # 'optim_state_dict': '/ncluster/runs/txl.09/optimizer-best.pt',
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e6,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git',
        'dataset': 'git',
    },
    'valid_custom': 'https://github-lm.s3.amazonaws.com/mailman.txt',
}

one_small_machine_git = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3.16xlarge',
    'local_batch_size': 6,
    'machines': 1,
    'large': True,
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e6,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git',
        'dataset': 'git',
    },
    'valid_custom': 'https://github-lm.s3.amazonaws.com/mailman.txt',
}

one_small_machine_git_checkpoint = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3.16xlarge',
    'local_batch_size': 6,
    'machines': 1,
    'large': True,
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e5,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git',
        'dataset': 'git',
    },
    'valid_custom': 'https://github-lm.s3.amazonaws.com/mailman.txt',
}

# 93.6% checkpoint + p3dn instance
one_small_machine_git_checkpoint_p3dn = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 18,
    'machines': 1,
    'large': True,
    'checkpoint_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-85-model.pt',
    'optimizer_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-85-optimizer.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e5,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git',
        'dataset': 'git',
    },
    'valid_custom': 'https://github-lm.s3.amazonaws.com/mailman.txt',
}

# 93.6% checkpoint + p3dn instance + long warmup
one_p3dn_machine_git_newcheckpoint = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 18,
    'machines': 1,
    'large': True,
    'checkpoint_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-85-model.pt',
    'optimizer_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-85-optimizer.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e6,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git',
        'dataset': 'git',
    },
    'valid_custom': 'https://github-lm.s3.amazonaws.com/mailman.txt',
}

one_p3_machine_biggit_newcheckpoint = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3.16xlarge',
    'local_batch_size': 6,
    'machines': 1,
    'large': True,
    'checkpoint_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-85-model.pt',
    'optimizer_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-85-optimizer.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e6,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git_85gb',
        'dataset': 'git',
    },
    'valid_custom': 'https://github-lm.s3.amazonaws.com/mailman_85gb.txt',
}

one_p3_machine_biggit_newcheckpoint_nodrop = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3.16xlarge',
    'local_batch_size': 6,
    'machines': 1,
    'large': True,
    'nodrop': True,
    'checkpoint_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-87-model.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e6,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git_85gb',
        'dataset': 'git',
    },
    'valid_custom': 'https://github-lm.s3.amazonaws.com/mailman_85gb.txt',
}

one_p3_machine_biggit_newcheckpoint_nodrop_p3dn = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 18,
    'machines': 1,
    'large': True,
    'nodrop': True,
    'checkpoint_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-87-model.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e6,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git_85gb',
        'dataset': 'git',
    },
    'valid_custom': 'https://github-lm.s3.amazonaws.com/mailman_85gb.txt',
}

# small model, 8x p3d
small_p3dn_8x = {
    'base_lr': 0.001 / 16,   # add 4x lr discount because 4x machine version didn't seem to work
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 18,
    'machines': 8,
    'large': True,
    'nodrop': True,
    'checkpoint_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-87-model.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e6,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git_85gb',
        'dataset': 'git',
    },
    'valid_custom': 'https://github-lm.s3.amazonaws.com/mailman_85gb.txt',
}

chimera = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 5,
    'machines': 1,
    'nodrop': True,
    'huge': True,
    'checkpoint_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-86-model.pt',
    'checkpoint_overwrite_secondary': 'https://s3.amazonaws.com/yaroslavvb2/data/git800-86-model.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e5,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git',
        'dataset': 'git',
    }
}

chimera_no = {
    'base_lr': 0.001 / 4 / 2,  # extra penalty of 2 because earlier run looked bad
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 5,
    'machines': 8,
    'nodrop': True,
    'huge': True,
    'checkpoint_overwrite':  'https://s3.amazonaws.com/yaroslavvb2/data/git800-86-model.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e5,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git_85gb',
        'dataset': 'git',
    }
}

# reduce warmup 5x, standard lr again
chimera_one = {
    'base_lr': 0.001 / 4,  # 
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 5,
    'machines': 1,
    'nodrop': True,
    'huge': True,
    'checkpoint_overwrite':  'https://s3.amazonaws.com/yaroslavvb2/data/git800-unknown-model.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 10e5,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git_85gb',
        'dataset': 'git',
    },
    'valid_custom': 'https://github-lm.s3.amazonaws.com/mailman_85gb.txt',
}

# 2x machines, 0.5x learning rate
chimera_two = {
    'base_lr': 0.001 / 4 / 2,
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 5,
    'machines': 2,
    'nodrop': True,
    'huge': True,
    'checkpoint_overwrite':  'https://s3.amazonaws.com/yaroslavvb2/data/git800-unknown-model.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 10e5,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git_85gb',
        'dataset': 'git',
    },
    'valid_custom': 'https://github-lm.s3.amazonaws.com/mailman_85gb.txt',
}


one_p3_machine_biggit_newcheckpoint_nodrop_p3dn_4x = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 18,
    'machines': 4,
    'large': True,
    'nodrop': True,
    'checkpoint_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-87-model.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e6,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git_85gb',
        'dataset': 'git',
    },
    'valid_custom': 'https://github-lm.s3.amazonaws.com/mailman_85gb.txt',
}

one_small_machine_newgit_checkpoint = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3.16xlarge',
    'local_batch_size': 6,
    'machines': 1,
    'large': True,
    'checkpoint_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-86-model.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e5,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git',
        'dataset': 'git',
    }
}

one_small_machine_newgit_checkpoint_nodrop = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3.16xlarge',
    'local_batch_size': 6,
    'machines': 1,
    'large': True,
    'checkpoint_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-86-model.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'nodrop': True,
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e5,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git',
        'dataset': 'git',
    }
}

# sanity check that p3dn's are not broken
one_small_machine_newgit_checkpoint_p0 = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 6,
    'machines': 1,
    'large': True,
    'checkpoint_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-86-model.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e5,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git',
        'dataset': 'git',
    }
}


one_small_machine_newgit_checkpoint_p3dn = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 18,
    'machines': 1,
    'large': True,
    'checkpoint_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-86-model.pt',
    'optimizer_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-86-optimizer.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e5,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git',
        'dataset': 'git',
    }
}

# 3x lower learning rate
one_small_machine_newgit_checkpoint_p3dn_slow = {
    'base_lr': 0.001 / 4 / 3,
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 18,
    'machines': 1,
    'large': True,
    'checkpoint_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-86-model.pt',
    'optimizer_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-86-optimizer.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e5,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git',
        'dataset': 'git',
    }
}


one_small_machine_newgit_checkpoint_p3dn2x = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 9,
    'machines': 2,
    'large': True,
    'checkpoint_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-86-model.pt',
    'optimizer_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-86-optimizer.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e5,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git',
        'dataset': 'git',
    }
}

one_small_machine_newgit_p3dn_huge = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 5,
    'machines': 1,
    'huge': True,
    'checkpoint_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-86-model.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e5,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git',
        'dataset': 'git',
    }
}

one_small_machine_newgit_p3dn_huge_nodrop = {
    'base_lr': 0.001 / 4,
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 5,
    'machines': 1,
    'huge': True,
    'checkpoint_overwrite': 'https://s3.amazonaws.com/yaroslavvb2/data/git360-86-model.pt',
    'checkpoint': 'github-projects_p3dn-2d_best.pt',  # us-east-1
    'nodrop': True,
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e5,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git',
        'dataset': 'git',
    }
}

one_machine_git = {
    'base_lr': 0.001 / 4,  # Divide by 2 to counteract batch adjustment
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 18,
    'machines': 1,
    'large': True,
    # 'checkpoint': '/ncluster/runs/txl.09/model-best.pt',  # us-east-1
    # 'optim_state_dict': '/ncluster/runs/txl.09/optimizer-best.pt',
    'extra_worker_params': {
        'fp16': True,
        'warmup_tokens': 50e6,
        'dynamic_loss_scale': True,
        'scheduler': 'constant',
        'data': 'data/git',
        'dataset': 'git',
    },
    'valid_custom': 'https://github-lm.s3.amazonaws.com/mailman.txt',
}

# Match https://github.com/kimiyoung/transformer-xl/blob/master/tf/scripts/wt103_large_tpu.sh
# Differences: fp16, lamb, 0 warmup, untie_r (doesn't exist in pytorch)
# logs: ben-large-lamb-slow
one_machine_fp16_large = {
    'base_lr': 0.001 / 4,  # Divide by 4 to counteract batch adjustment
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 16,
    'machines': 1,
    'large': True,
}

# fork of one_machine_fp16_large
four_machine_fp16_large = {
    'base_lr': 0.001 / 4,  # Divide by 4 to counteract batch adjustment
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 16,
    'machines': 4,
    'large': True,
}

# logs: ben-eight.01
eight_machine_fp16_large = {
    'base_lr': 0.001 / 4,  # Divide by 4 to counteract batch adjustment
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 16,
    'machines': 8,
    'large': True,
}

# /ncluster/runs.new/yaro-fp16.09
one_machine_fp16_2xlr = {
    'base_lr': 0.000125 * 5 / 3 * 2,  # from ben-big-lr.09
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 96,
    'machines': 1,
    'extra_worker_params': {
        'fp16': True,
        'dynamic_loss_scale': True,
    }
}

# /ncluster/runs.new/yaro-fp16.10
one_machine_fp16_4xlr = {
    'base_lr': 0.000125 * 5 / 3 * 4,  # from ben-big-lr.09
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 96,
    'machines': 1,
    'extra_worker_params': {
        'fp16': True,
        'dynamic_loss_scale': True,
    }
}

# /ncluster/runs.new/yaro-fp16.11
# nans after 30 seconds
one_machine_fp16_8xlr = {
    'base_lr': 0.000125 * 5 / 3 * 8,  # from ben-big-lr.09
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 96,
    'machines': 1,
    'extra_worker_params': {
        'fp16': True,
        'dynamic_loss_scale': True,
    }
}

# smaller p3.16 machine, logs: ben-bpe
one_machine_fp16_small = {
    'base_lr': 0.000125 * 5 / 3 / 2,  # from ben-big-lr.09
    'instance_type': 'p3.16xlarge',
    'local_batch_size': 96 // 2,
    'machines': 1,
    'extra_worker_params': {
        'fp16': True,
        'dynamic_loss_scale': True,
    }
}

one_machine_fp16_checkpoint = {
    'base_lr': 0.000125 * 5 / 3,  # from ben-big-lr.09
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 96,
    'machines': 1,
    'checkpoint': '/ncluster/runs.new/yaro-one.08/model-1.pt',
    'extra_worker_params': {
        'fp16': True,
        'dynamic_loss_scale': True,
    }
}

# /ncluster/runs.new/yaro-two-fp16.04 (with checkpoints)
two_machines_fp16 = {
    'base_lr': 0.000125 * 5 / 3,  # from ben-big-lr.09
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 96,
    'machines': 2,
    'extra_worker_params': {
        'fp16': True,
        'dynamic_loss_scale': True,
    }
}

two_machines = {
    'base_lr': 0.000125 * 5 / 3,  # yaro-two.07
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 96,
    'machines': 2,
}

# yaro-four
four_machines = {
    'base_lr': 0.000125,  # remove ben's 5/3 tweak, and additional penalty of 2x
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 96,
    'machines': 4,
    'extra_worker_params': {
        'fp16': True,
        'dynamic_loss_scale': True,
    }
}

eight_machines = {
    'base_lr': 0.000125 / 2,  # remove ben's 5/3 tweak, and additional penalty of 2x
    'instance_type': 'p3dn.24xlarge',
    'local_batch_size': 96,
    'machines': 8,
    'checkpoint': '/ncluster/runs.new/yaro-one.08/model-1.pt',
    'extra_worker_params': {
        'fp16': True,
        'dynamic_loss_scale': True,
        'warmup_tokens': 50e6,
    }
}


def dict_to_args(dict_: dict):
    def item_to_arg(item: tuple):
        k, v = item
        if v is False or v is None:
            return ''
        if v is True:
            return f'--{k}'
        return f'--{k} {v}'

    return ' '.join([item_to_arg(item) for item in dict_.items()])


HUGE_ARGS = {
    'n_layer': 18*3,
    'd_model': 1024,
    'n_head': 16,
    'd_head': 64,
    'd_inner': 4096,
    'dropout': 0.2,
    'dropatt': 0.2,
    'optim': 'lamb',
    'warmup_tokens': 0,
    'tgt_len': 384,
    'mem_len': 384,
    'eval_tgt_len': 128,
    'fp16': True,
    'dynamic_loss_scale': True,
    'init_std': 0.005,
    'div_val': 4,
    'freeze_below': 18,
}


# Match https://github.com/kimiyoung/transformer-xl/blob/master/tf/scripts/wt103_large_tpu.sh
LARGE_ARGS = {
    'n_layer': 18,
    'd_model': 1024,
    'n_head': 16,
    'd_head': 64,
    'd_inner': 4096,
    'dropout': 0.2,
    'dropatt': 0.2,
    'optim': 'lamb',
    'warmup_tokens': 0,
    'tgt_len': 384,
    'mem_len': 384,
    'eval_tgt_len': 128,
    'fp16': True,
    'dynamic_loss_scale': True,
    'init_std': 0.005,
    'div_val': 4,
}

# Roughly match https://github.com/kimiyoung/transformer-xl/blob/master/pytorch/run_wt103_base.sh
# Divisible by 8 for fp16 compatibility.
SMALL_ARGS = {
    'n_layer': 16,
    'd_model': 512,
    'n_head': 8,
    'd_head': 48,
    'd_inner': 2048,
    'dropout': 0.1,
    'dropatt': 0.0,
    'optim': 'lamb',
    'tgt_len': 128,
    'mem_len': 128,
    'eval_tgt_len': 128,
}


def _get_nccl_params():
    params = f'NCCL_DEBUG=VERSION '

    params += f'NCCL_MIN_NRINGS={args.num_rings} ' \
              f'NCCL_MAX_NRINGS={args.num_rings} '
    return params


def main():
    ncluster.set_backend('aws')

    if args.config:
        assert not args.instance_type, "specify instance_type as part of config"
        assert not args.machines, "specify number of machines as part of config"
        assert re.match('\\w+', args.config)
        assert args.config in globals(), f'no config called {args.config}'
        config = eval(args.config)

    else:  # setting config vars through command-line flags
        assert args.instance_type
        assert args.machines
        config = {'base_lr': 0.000125 * 5 / 3,
                  'local_batch_size': 96,
                  'instance_type': args.instance_type,
                  'machines': args.machines}

    config = AttrDefault(str, config)  # easier access to dictionary entries
    config.image_name = IMAGE_NAME
    config.conda_env = CONDA_ENV

    if args.conda_env:
        config.conda_env = args.conda_env
        print("Using non-standard conda env ", config.conda_env)
    if args.image_name:
        config.image_name = args.image_name
        print("Using non-standard image ", config.image_name)

    instance_info = ncluster.aws_backend.INSTANCE_INFO[config.instance_type]
    num_gpus_per_machine = instance_info['gpus']

    job = ncluster.make_job(name=args.name,
                            run_name=f"{args.name}",
                            num_tasks=config.machines,
                            image_name=config.image_name,
                            instance_type=config.instance_type,
                            spot=not args.nospot,
                            skip_setup=args.skip_setup)

    job.rsync('.')
    job.run(f'killall python || echo failed && '  # kill previous run
            f'source activate {config.conda_env} && ' +
            # protobuf https://github.com/tensorflow/models/issues/3995
            f'pip uninstall -y protobuf && ' +
            f'pip install -U protobuf && ' +
            f'pip install -r requirements.txt')
    # job.run('bash get_git_data.sh')
    # job.run('bash get_git_data_85gb.sh')
    # job.run('curl "https://github-lm.s3.amazonaws.com/github-projects_p3dn-2d_best.pt" '
    #         '-o github-projects_p3dn-2d_best.pt')

    local_batch_size = config.local_batch_size
    base_lr = config.base_lr

    num_workers = num_gpus_per_machine * config.machines
    global_batch_size = local_batch_size * num_workers
    print("using global batch ", global_batch_size)  # 512=8*32*2*1

    # linear LR scaling (https://arxiv.org/abs/1706.02677)
    lr = base_lr * (global_batch_size / BASE_LR_BATCHSIZE)

    # worker parameters with training setup
    worker_params = {
        'seed': 1111,
        'data': 'data/wikitext-103',
        'dataset': 'wt103',
        'adaptive': True,
        'log_interval': 100,
        'eval_interval': 500,
        'max_tokens': int(1.5e9),
        'logdir': job.logdir,
        'lr': lr,
        'batch_size': local_batch_size,
        'eta_min': lr / 10,
        'run_name': args.name,
    }

    if config.huge:
        worker_params.update(HUGE_ARGS)
    elif config.large:
        worker_params.update(LARGE_ARGS)
    else:
        worker_params.update(SMALL_ARGS)

    if config.nodrop:
        #'dropout': 0.2,
        #'dropatt': 0.2,
        worker_params['dropout'] = 0
        worker_params['dropatt'] = 0

    if config.huge:
        worker_params['log_interval'] = 1
        worker_params['eval_interval'] = 100

    user_params = {}
    # pass through some user-provided settings that were arguments to the launcher script
    if args.checkpoint_each_epoch:
        user_params['checkpoint_each_epoch'] = args.checkpoint_each_epoch
    if config.warmup_tokens:
        user_params['warmup_tokens'] = config.warmup_tokens
    if args.skip_files:
        user_params['skip_files'] = args.skip_files

    if args.checkpoint or config.checkpoint:
        user_params['checkpoint'] = util.one_of([args.checkpoint, config.checkpoint])
    if args.optim_state_dict or config.optim_state_dict:
        user_params['optim_state_dict'] = util.one_of([args.optim_state_dict, config.optim_state_dict])

    if config.valid_custom:
        downloaded_fn = util.download_from_s3(config.valid_custom, job)
        user_params['valid_custom'] = downloaded_fn
    if args.valid_custom:
        downloaded_fn = util.download_from_s3(config.valid_custom, job)
        user_params['valid_custom'] = downloaded_fn

    if args.wiki:
        worker_params.update({
            'data': 'data/wikiextracted',
            'dataset': 'wiki',
            'dropatt': 0.1,
            'dropout': 0.1,
        })
    if args.git:
        worker_params.update({
            'data': 'data/git',
            'dataset': 'git',
            'div_val': 1,
            'bpe': True,
            'adaptive': False,
        })

    if args.bpe:
        worker_params.update({
            'div_val': 1,
            'bpe': True,
            'adaptive': False,
        })

    if config.checkpoint_overwrite:
        downloaded_fn = util.download_from_s3(config.checkpoint_overwrite, job)
        print(f"downloaded {downloaded_fn} from {config.checkpoint_overwrite}")
        user_params["checkpoint"] = downloaded_fn
    if config.checkpoint_overwrite_secondary:
        downloaded_fn = util.download_from_s3(config.checkpoint_overwrite_secondary, job)
        print(f"downloaded secondary {downloaded_fn} from {config.checkpoint_overwrite_secondary}")
        user_params["checkpoint_secondary"] = downloaded_fn
    if config.optimizer_overwrite:
        downloaded_fn = util.download_from_s3(config.optimizer_overwrite, job)
        user_params["optim_state_dict"] = downloaded_fn

    worker_params.update(user_params)

    if config.extra_worker_params:
        worker_params.update(config.extra_worker_params)

    nccl_params = _get_nccl_params()

    for i, task in enumerate(job.tasks):
        dist_params = \
            f'--nproc_per_node={num_gpus_per_machine} ' \
            f'--nnodes={config.machines} --node_rank={i} ' \
            f'--master_addr={job.tasks[0].ip} --master_port={6016}'
        cmd = f'{nccl_params} python -m torch.distributed.launch {dist_params} train.py {dict_to_args(worker_params)}'
        task.run(f'echo {cmd} > {job.logdir}/task-{i}.cmd')  # save command-line
        task.run(cmd, non_blocking=True)

    print(f"Logging to {job.logdir}")


if __name__ == '__main__':
    main()