

AutoMoE: Neural Architecture Search for Efficient Sparsely Activated Transformers
ICLR 2023 Submission


Quick start of AutoMoE

1) Get preprocessed data
bash configs/wmt14.en-de/get_preprocessed.sh
bash configs/wmt14.en-fr/get_preprocessed.sh
bash configs/wmt19.en-de/get_preprocessed.sh

2) Train supernet
python -B train.py --configs=configs/wmt19.en-de/supertransformer/space0.yml --save-dir /moe_6e/supernet --encoder-n-experts 1 2 3 4 5 6 --decoder-n-experts 1 2 3 4 5 6 --encoder-each-expert-ffn-dim 0 1 --decoder-each-expert-ffn-dim 0 1 --no-epoch-checkpoints --max-tokens 8192 --max-update 40000 --update-freq 16 --num_workers 1

3) Generate latency dataset (optional, if you want to use latency predictor)
CUDA_VISIBLE_DEVICES=0 python latency_dataset.py --configs=configs/wmt19.en-de/latency_dataset/gpu_titanxp.yml --encoder-n-experts 1 2 3 4 5 6 --decoder-n-experts 1 2 3 4 5 6 --encoder-each-expert-ffn-dim 0 1 --decoder-each-expert-ffn-dim 0 1 --lat-dataset-path /moe_6e/genlatdata/wmt19.en-de_gpu_titanxp.csv --max-tokens 8192 --max-update 40000 --update-freq 16 

3) Train latency predictor (optional)
python latency_predictor.py --configs=configs/wmt19.en-de/latency_predictor/gpu_titanxp.yml --feature-norm 640 6 2048 6 640 6 2048 6 6 2 6 6 3072 3072 --feature-dim 14 --lat-dataset-path /moe_6e/genlatdata/wmt19.en-de_gpu_titanxp.csv --ckpt-path /moe_6e/latpred/wmt19.en-de_gpu_titanxp.pt 

4) Run evolutionary search
CUDA_VISIBLE_DEVICES=0 python evo_search.py --configs=configs/wmt19.en-de/supertransformer/space0.yml --evo-configs=configs/wmt19.en-de/evo_search/wmt19ende_titanxp.yml --restore-file /moe_6e/supernet/checkpoint_best.pt --ckpt-path /moe_6e/latpred/wmt19.en-de_gpu_titanxp.pt --feature-norm 640 6 2048 6 640 6 2048 6 6 2 6 6 3072 3072 --write-config-path /moe_6e/evosearch/wmt19.en-de_gpu_titanxp.yml --encoder-n-experts 1 2 3 4 5 6 --decoder-n-experts 1 2 3 4 5 6 --encoder-each-expert-ffn-dim 0 1 --decoder-each-expert-ffn-dim 0 1 --max-tokens 8192 --max-update 40000 --update-freq 16 --validation-metric loss --latency-constraint 200 --latency-compute gold --evo-iter 10 --latiter 100

5) Get inherited performance of best subnet
bash configs/wmt19.en-de/test.sh /moe_6e/supernet/checkpoint_best.pt /moe_6e/evosearch/wmt19.en-de_gpu_titanxp.yml /moe_6e/perf_inherited_test normal 0 valid >> /moe_6e/perf_inherited_test/wmt19.en-de_gpu_titanxp.txt
bash configs/wmt19.en-de/test.sh /moe_6e/supernet/checkpoint_best.pt /moe_6e/evosearch/wmt19.en-de_gpu_titanxp.yml /moe_6e/perf_inherited_test normal >> /moe_6e/perf_inherited_test/wmt19.en-de_gpu_titanxp.txt

6) Get gold latency of best subnet
CUDA_VISIBLE_DEVICES=0 python train.py --configs=/moe_6e/evosearch/wmt19.en-de_gpu_titanxp.yml --sub-configs=configs/wmt19.en-de/subtransformer/common.yml --latgpu --save-dir /moe_6e/latency_flops_inherited_test --max-tokens 8192 --max-update 40000 --update-freq 16 --num_workers 1 >> /moe_6e/latency_flops_inherited_test/wmt19.en-de_gpu_titanxp.txt

7) Get FLOPs of best subnet
CUDA_VISIBLE_DEVICES=0 python train.py --configs=/moe_6e/evosearch/wmt19.en-de_gpu_titanxp.yml --sub-configs=configs/wmt19.en-de/subtransformer/common.yml --profile-flops --save-dir /moe_6e/latency_flops_inherited_test --max-tokens 8192 --max-update 40000 --update-freq 16 --num_workers 1 >> /moe_6e/latency_flops_inherited_test/wmt19.en-de_gpu_titanxp.txt

8) Train best subnet from scratch
python -B train.py --configs=/moe_6e/evosearch/wmt19.en-de_gpu_titanxp.yml --save-dir /moe_6e/effnet --sub-configs=configs/wmt19.en-de/subtransformer/common.yml --no-epoch-checkpoints --max-tokens 8192 --max-update 40000 --update-freq 16 > /moe_6e/effnet/run.out

9) Get scratch performance of best subnet
bash configs/wmt19.en-de/test.sh /moe_6e/effnet/checkpoint_best.pt /moe_6e/evosearch/wmt19.en-de_gpu_titanxp.yml /moe_6e/perf_scratch_test normal 0 valid >> /moe_6e/perf_scratch_test/wmt19.en-de_gpu_titanxp.txt
bash configs/wmt19.en-de/test.sh /moe_6e/effnet/checkpoint_best.pt /moe_6e/evosearch/wmt19.en-de_gpu_titanxp.yml /moe_6e/perf_scratch_test normal >> /moe_6e/perf_scratch_test/wmt19.en-de_gpu_titanxp.txt



