Fast implementation of BERT inference directly on NVIDIA (CUDA, CUBLAS) and Intel MKL
Highly customized and optimized BERT inference directly on NVIDIA (CUDA, CUBLAS) or Intel MKL, without tensorflow and its framework overhead.
ONLY BERT (Transformer) is supported.
- Tesla P4
- 28 * Intel(R) Xeon(R) CPU E5-2680 v4 @ 2.40GHz
- Debian GNU/Linux 8 (jessie)
- gcc (Debian 4.9.2-10+deb8u1) 4.9.2
- CUDA: release 9.0, V9.0.176
- MKL: 2019.0.1.20181227
- tensorflow: 1.12.0
- BERT: seq_length = 32
|batch size||128 (ms)||32 (ms)|
|batch size||128 (ms)||1 (ms)|
Note: MKL should be run under
OMP_NUM_THREADS=? to control its thread
number. Other environment variables and their possible values includes:
cuBERT can be accelerated by Tensor Core and Mixed Precision on NVIDIA Volta and Turing GPUs. We support mixed precision as variables stored in fp16 with computation taken in fp32. The typical accuracy error is less than 1% compared with single precision inference, while the speed achieves more than 2x acceleration.
We support following 2 pooling method.
- The standard BERT pooler, which is defined as:
with tf.variable_scope("pooler"): # We "pool" the model by simply taking the hidden state corresponding # to the first token. We assume that this has been pre-trained first_token_tensor = tf.squeeze(self.sequence_output[:, 0:1, :], axis=1) self.pooled_output = tf.layers.dense( first_token_tensor, config.hidden_size, activation=tf.tanh, kernel_initializer=create_initializer(config.initializer_range))
- Simple average pooler:
self.pooled_output = tf.reduce_mean(self.sequence_output, axis=1)
Following outputs are supported:
Build from Source
mkdir build && cd build # if build with CUDA cmake -DCMAKE_BUILD_TYPE=Release -DcuBERT_ENABLE_GPU=ON -DCUDA_ARCH_NAME=Common .. # or build with MKL cmake -DCMAKE_BUILD_TYPE=Release -DcuBERT_ENABLE_MKL_SUPPORT=ON .. make -j4 # install to /usr/local # it will also install MKL if -DcuBERT_ENABLE_MKL_SUPPORT=ON sudo make install
If you would like to run tfBERT_benchmark for performance comparison, please first install tensorflow C API from https://www.tensorflow.org/install/lang_c.
Run Unit Test
Download BERT test model
and put them under dir
build before run
make test or
We provide simple Python wrapper by Cython, and it can be built and installed after C++ building as follows:
cd python python setup.py bdist_wheel # install pip install dist/cuBERT-xxx.whl # test python cuBERT_test.py
Please check the Python API usage and examples at cuBERT_test.py for more details.
Java wrapper is implemented through JNA . After installing maven and C++ building, it can be built as follows:
cd java mvn clean package # -DskipTests
When using Java JAR, you need to specify
jna.library.path to the
libcuBERT.so if it is not installed to the system path.
jna.encoding should be set to UTF8 as
in the JVM start-up script.
Please check the Java API usage and example at ModelTest.java for more details.
Pre-built python binary package (currently only with MKL on Linux) can be installed as follows:
Download and install MKL to system path.
Download the wheel package and
pip install cuBERT-xxx-linux_x86_64.whl
python -c 'import libcubert'to verify your installation.
cuBERT is built with protobuf-c to avoid version and code conflicting with tensorflow protobuf.
Libraries compiled by CUDA with different versions are not compatible.
MKL is dynamically linked. We install both cuBERT and MKL in
sudo make install.
We assume the typical usage case of cuBERT is for online serving, where concurrent requests of different batch_size should be served as fast as possible. Thus, throughput and latency should be balanced, especially in pure CPU environment.
As the vanilla class Bert is not thread-safe
because of its internal buffers for computation, a wrapper class BertM
is written to hold locks of different
Bert instances for thread safety.
BertM will choose one underlying
Bert instance by a round-robin
manner, and consequence requests of the same
Bert instance might be
queued by its corresponding lock.
Bert is placed on one GPU card. The maximum concurrent requests is
the number of usable GPU cards on one machine, which can be controlled
CUDA_VISIBLE_DEVICES if it is specified.
For pure CPU environment, it is more complicate than GPU. There are 2 level of parallelism:
Request level. Concurrent requests will compete CPU resource if the online server itself is multi-threaded. If the server is single-threaded (for example some server implementation in Python), things will be much easier.
Operation level. The matrix operations are parallelized by OpenMP and MKL. The maximum parallelism is controlled by
MKL_NUM_THREADS, and many other environment variables. We refer our users to first read Using Threaded Intel® MKL in Multi-Thread Application and Recommended settings for calling Intel MKL routines from multi-threaded applications .
Thus, we introduce
CUBERT_NUM_CPU_MODELS for better control of request
level parallelism. This variable specifies the number of
created on CPU/memory, which acts same like
If you have limited number of CPU cores (old or desktop CPUs, or in Docker), it is not necessary to use
CUBERT_NUM_CPU_MODELS. For example 4 CPU cores, a request-level parallelism of 1 and operation-level parallelism of 4 should work quite well.
But if you have many CPU cores like 40, it might be better to try with request-level parallelism of 5 and operation-level parallelism of 8.
MKL_NUM_THREADS defines how many threads
one model could use, and
CUBERT_NUM_CPU_MODELS defines how many models in
Again, the per request latency and overall throughput should be balanced,
and it diffs from model
batch_size, your CPU cores, your
server QPS, and many many other things. You should take a lot benchmark
to achieve the best trade-off. Good luck!