AutoMoE: Neural Architecture Search for Efficient Sparsely Activated Transformers

This repository contains code, data and pretrained models used in AutoMoE (pre-print). This repository builds on Hardware Aware Transformer (HAT)'s repository.

AutoMoE Framework

AutoMoE Key Result

The following table shows the performance of AutoMoE vs. baselines on standard machine translation benchmarks: WMT'14 En-De, WMT'14 En-Fr and WMT'19 En-De.

WMT’14 En-De	Network	# Active Params (M)	Sparsity (%)	FLOPs (G)	BLEU	GPU Hours
Transformer	Dense	176	0	10.6	28.4	184
Evolved Transformer	NAS over Dense	47	0	2.9	28.2	2,192,000
HAT	NAS over Dense	56	0	3.5	28.2	264
AutoMoE (6 Experts)	NAS over Sparse	45	62	2.9	28.2	224

WMT’14 En-Fr	Network	# Active Params (M)	Sparsity (%)	FLOPs (G)	BLEU	GPU Hours
Transformer	Dense	176	0	10.6	41.2	240
Evolved Transformer	NAS over Dense	175	0	10.8	41.3	2,192,000
HAT	NAS over Dense	57	0	3.6	41.5	248
AutoMoE (6 Experts)	NAS over Sparse	46	72	2.9	41.6	236
AutoMoE (16 Experts)	NAS over Sparse	135	65	3.0	41.9	236

WMT’19 En-De	Network	# Active Params (M)	Sparsity (%)	FLOPs (G)	BLEU	GPU Hours
Transformer	Dense	176	0	10.6	46.1	184
HAT	NAS over Dense	63	0	4.1	45.8	264
AutoMoE (2 Experts)	NAS over Sparse	45	41	2.8	45.5	248
AutoMoE (16 Experts)	NAS over Sparse	69	81	3.2	45.9	248

Quick Setup

(1) Install

Run the following commands to install AutoMoE:

git clone https://github.com/UBC-NLP/AutoMoE.git
cd AutoMoE
pip install --editable .

(2) Prepare Data

Run the following commands to download preprocessed MT data:

bash configs/[task_name]/get_preprocessed.sh

where [task_name] can be wmt14.en-de or wmt14.en-fr or wmt19.en-de.

(3) Run full AutoMoE pipeline

Run the following commands to start AutoMoE pipeline:

python generate_script.py --task wmt14.en-de --output_dir /tmp --num_gpus 4 --trial_run 0 --hardware_spec gpu_titanxp --max_experts 6 --frac_experts 1 > automoe.sh
bash automoe.sh

where,

• task - MT dataset to use: wmt14.en-de or wmt14.en-fr or wmt19.en-de (default: wmt14.en-de)
• output_dir - Output directory to write files generated during experiment (default: /tmp)
• num_gpus - Number of GPUs to use (default: 4)
• trial_run - Run trial run (useful to quickly check if everything runs fine without errors.): 0 (final run), 1 (dry/dummy/trial run) (default: 0)
• hardware_spec - Hardware specification: gpu_titanxp (For GPU) (default: gpu_titanxp)
• max_experts - Maximum experts (for Supernet) to use (default: 6)
• frac_experts - Fractional (varying FFN. intermediate size) experts: 0 (Standard experts) or 1 (Fractional) (default: 1)
• supernet_ckpt - Skip supernet training by specifiying checkpoint from pretrained models (default: None)
• latency_compute - Use (partially) gold or predictor latency (default: gold)
• latiter - Number of latency measurements for using (partially) gold latency (default: 100)
• latency_constraint - Latency constraint in terms of milliseconds (default: 200)
• evo_iter - Number of iterations for evolutionary search (default: 10)

Contact

If you have questions, contact Ganesh (ganeshjwhr@gmail.com), Subho (Subhabrata.Mukherjee@microsoft.com) and/or create GitHub issue.

Citation

If you use this code, please cite:

@misc{jawahar2022automoe,
      title={AutoMoE: Neural Architecture Search for Efficient Sparsely Activated Transformers}, 
      author={Ganesh Jawahar and Subhabrata Mukherjee and Xiaodong Liu and Young Jin Kim and Muhammad Abdul-Mageed and Laks V. S. Lakshmanan and Ahmed Hassan Awadallah and Sebastien Bubeck and Jianfeng Gao},
      year={2022},
      eprint={2210.07535},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}

License

See LICENSE.txt for license information.

Acknowledgements

• Hardware Aware Transformer from mit-han-lab
• fairseq from facebookresearch

Trademarks

This project may contain trademarks or logos for projects, products, or services. Authorized use of Microsoft trademarks or logos is subject to and must follow Microsoft's Trademark & Brand Guidelines. Use of Microsoft trademarks or logos in modified versions of this project must not cause confusion or imply Microsoft sponsorship. Any use of third-party trademarks or logos are subject to those third-party's policies.