โก
Substation: Optimized Transformers Substation is a project to optimize transformers using data movement analysis.
This code is presently at a research-and-development stage. We are actively working to make it both faster and more usable.
For more background, please see our paper, Data Movement Is All You Need: A Case Study on Optimizing Transformers. If you use our code, please cite the paper:
@article{ivanov2020data,
title={Data Movement Is All You Need: A Case Study on Optimizing Transformers},
author={Ivanov, Andrei and Dryden, Nikoli and Ben-Nun, Tal and Li, Shigang and Hoefler, Torsten},
journal={arXiv preprint arXiv:2007.00072},
year={2020}
}
Current Performance
We presently include configurations for two versions of a single BERT-large encoder layer:
- Batch size 8 and max sequence length 512.
- Batch size 96 and max sequence length 128.
These benchmarks were run on the Lassen supercomputer. Note that the Nvidia V100s this system uses are the SXM2 variety, with a peak of 125 tflop/s using Tensor Cores. We compare with the same transformer architecture implemented in TensorFlow (with XLA), PyTorch, and DeepSpeed. These results are with the latest version of our code, but see our paper for other details.
All times are in milliseconds (ms).
BERT-large, batch size 8, max sequence length 512 runtime
PyTorch | TensorFlow+XLA | DeepSpeed | Substation |
---|---|---|---|
9.14 | 8.4 | 7.6 | 6.71 |
BERT-large, batch size 96, max sequence length 128 runtime
PyTorch | TensorFlow+XLA | DeepSpeed | Substation |
---|---|---|---|
18.43 | n/a | 16.19 | 15.42 |
Usage
Note: We are actively working to improve the usability for standard deep learning workflows.
Our encoder implementation is available as a PyTorch module in pytorch_module/encoder.py
. Whenever you create a Substation encoder, you must specify an associated set of layouts and other configurations (see below for generating one yourself). We have provided the configurations used for the two BERT-large versions above as layouts-bert-b8-l512-h16-e1024.pickle
and layouts-bert-b96-l128-h16-e1024.pickle
, respectively. These configurations are optimized for the specific configuration and hardware, but should run for other problem sizes and on other hardware. The underlying optimized implementation for the encoder will be generated and compiled the first time you use it.
For performance benchmarking, we provide the run_encoder.py
script. See its --help
information for details.
Generating New Configurations
If you want to get the best performance for your particular problem configuration and/or hardware, you will need to generate a configuration. This involves two phases: benchmarking to gather performance data, then configuration selection.
Benchmarking
Warning: This can take a long time.
This exhaustively benchmarks the possible layouts (and other options) for every operator used in the encoder layer. There are two sets of benchmarks, one for tensor contractions (which uses cuBLAS) and one for our custom fused kernel implementations.
Tensor Contractions
These are located in tc_profiling
.
- Run
compile.sh
to build cuBLAS benchmarks. - Run
einsum_perms.py
(e.g.,einsum_perms.py --b 8 --j 512 --h 16 --i 1024
) to generate the benchmark configurations for each operator. - These configurations can be run with
runprof.py <config file>
.
Fused Kernels
These are run with the pytorch_module/benchmark.py
script. You specify the kernel to benchmark with --kernel name
. By default, this uses the batch size 8, sequence length 512 configuration of BERT-large. You can change the size using the --size
argument. For example:
python benchmark.py --kernel softmax --size "H=16,B=96,J=128,K=128,U=4096,N=1024,P=64"
See its --help
for more arguments.
You will need to run every tensor contraction and kernel benchmark.
Configuration Selection
These scripts are located in the config_selection
directory. First, collect the benchmark data into a directory. You can just copy the kernel benchmark output. Use the parse_tc_results.py
script to assemble the tensor contraction results and then copy them into the same directory.
Final configuration selection can then be run with python optimize.py --output_config my_layouts.pickle results-dir
.
Advanced
The optimize.py
script can use several strategies for performing configuration selection, controlled with the --graph_order
argument. The default, bp_first
, will optimize the encoder layer's backpropagation pass first, and then its forward pass. fp_first
will optimize forward propagation first, then backpropagation. bp_first
typically results in configurations that are faster than fp_first
. The third option, combined
, will optimize over forward and backpropagation simultaneously, and typically results in the fastest configurations. However, this approach is somewhat finnicky, and can often fail to find a valid layout. This can be worked around by telling the optimizer to "split" at certain variables using the --split_vars
argument.
The layouts-bert-b8-l512-h16-e1024.pickle
configuration was generated using optimize.py --graph_order combined --split_vars X LN1 LN2 LIN2 DLIN2
. The layouts-bert-b96-l128-h16-e1024.pickle
configuration was generated using optimize.py --graph_order combined --split_vars X DROP2 LN1
.
Contributors
This project is led by the Scalable Parallel Computing Lab at ETH Zurich.
See also the list of contributors.