Profiling GPT2 Inference Latency (FP32)

I ask the question here: does a GPT-2 Large forward pass fully utilize the hardware capabilities of an NVIDIA A100? Does this expensive hardware pay off for running Transformer inference? Large Language Models (LLMs) like GPT-2 have generated significant activity in optimizing machine learning inference along two major axes: cost of inference (largely driven by hardware or cloud cost), and latency of inference. The default choice for running inference has mostly been NVIDIA GPUs which are among the most expensive hardware components available for the task. Since researchers are investigating transformer architectures now to complement or even replace traditional architectures like convolutional networks I think it’s worthwhile to briefly consider what the performance characteristics of inference on these are like, to serve as comparison for the LLM.

At small batch sizes, some convolutional neural networks underutilize the compute capabilities of high-end GPUs but still benefit in latency from high memory bandwidth. Small kernels with low occupancy can bottleneck the critical path, resulting in bursts of compute intensity followed by low intensity. For example, the trace below shows the stream of kernels (zoomed out too far to read individual names) and some high-level performance metrics such as occupancy, compute, and memory throughput.

Inference of efficientnet_v2_l model with batch=1 on NVIDIA A100 (TRT engine)

Thus this brings a fair question on the tradeoff between cost and latency. If one expects to mostly do inference at small batch sizes (which is a common use case), could we perhaps use a less capable but cheaper machine to deliver the same results? We can’t answer this in a vacuum and need to know about the application, but for small batch inference of a particular convolutional network, a high-end GPU may not deliver value proportional to its increased cost. That depends on the model, software stack, latency target, and pricing, so I do not assume it in advance here. But what about LLMs?

LLMs change the game in many respects compared to their predecessors. Their models tend to have many more parameters than convolutional neural networks, so much that they often will not fully fit within the memory of a single GPU. Since GPU memory capacity tends to only increase significantly on the much higher-end GPUs this has pushed LLM inference into those GPUs, driving up costs. Transformer architectures mark another departure and important factor for evaluating performance. Convolutional layers themselves need to process relatively little amount of data (the convolution kernel,input, and output) and any data increase usually comes from increasing number of channels from the convolutional layer (a linear dependence). Multi-headed attention, on the other hand, computes comparatively large matrix multiplications involving query, key, and value matrices. The attention score matrix and its associated arithmetic grow quadratically with sequence length, while the projections and model weights scale linearly with the number of tokens. Is it possible then that LLMs may see kernels with greater occupancy or greater compute intensity and fewer critical path bottlenecks of low intensity such as seen with convolutional networks, leading to better use high-end GPUs? I investigate below.

GPT-2 Small Batch inference on NVIDIA A100

Basic Timing

I benchmark one forward pass with batch size 1 and sequence length 1024. The export disables the key/value cache and returns logits for every position, so this is full-sequence inference rather than autoregressive token generation.

I used a prebaked nvidia image on GCP with an nvidia A100 - screengrab of details below.

First I generated a TensorRT engine from this model and did the high-level benchmark (no hardware measurements, just times):

[12/04/2022-02:01:36] [I] === Trace details ===
[12/04/2022-02:01:36] [I] Trace averages of 10 runs:
[12/04/2022-02:01:36] [I] Average on 10 runs - GPU latency: 31.0561 ms - Host latency: 41.2429 ms (enqueue 0.151791 ms)
[12/04/2022-02:01:36] [I] Average on 10 runs - GPU latency: 31.0865 ms - Host latency: 41.274 ms (enqueue 0.148151 ms)
[12/04/2022-02:01:36] [I] Average on 10 runs - GPU latency: 31.0933 ms - Host latency: 41.2813 ms (enqueue 0.149274 ms)
[12/04/2022-02:01:36] [I] Average on 10 runs - GPU latency: 31.0886 ms - Host latency: 41.2748 ms (enqueue 0.137341 ms)
[12/04/2022-02:01:36] [I] Average on 10 runs - GPU latency: 31.0964 ms - Host latency: 41.2835 ms (enqueue 0.13833 ms)
[12/04/2022-02:01:36] [I] Average on 10 runs - GPU latency: 31.0831 ms - Host latency: 41.2698 ms (enqueue 0.155908 ms)
[12/04/2022-02:01:36] [I] Average on 10 runs - GPU latency: 31.086 ms - Host latency: 41.2733 ms (enqueue 0.140601 ms)
[12/04/2022-02:01:36] [I] Average on 10 runs - GPU latency: 31.0899 ms - Host latency: 41.2763 ms (enqueue 0.148071 ms)
[12/04/2022-02:01:36] [I] Average on 10 runs - GPU latency: 31.0868 ms - Host latency: 41.2747 ms (enqueue 0.151758 ms)
[12/04/2022-02:01:36] [I] 
[12/04/2022-02:01:36] [I] === Performance summary ===
[12/04/2022-02:01:36] [I] Throughput: 31.8446 qps
[12/04/2022-02:01:36] [I] Latency: min = 41.1393 ms, max = 41.3042 ms, mean = 41.273 ms, median = 41.2759 ms, percentile(90%) = 41.2895 ms, percentile(95%) = 41.2952 ms, percentile(99%) = 41.3042 ms
[12/04/2022-02:01:36] [I] Enqueue Time: min = 0.120361 ms, max = 0.22583 ms, mean = 0.148609 ms, median = 0.146973 ms, percentile(90%) = 0.173096 ms, percentile(95%) = 0.180908 ms, percentile(99%) = 0.22583 ms
[12/04/2022-02:01:36] [I] H2D Latency: min = 0.00842285 ms, max = 0.0283203 ms, mean = 0.0102436 ms, median = 0.00927734 ms, percentile(90%) = 0.012207 ms, percentile(95%) = 0.0171509 ms, percentile(99%) = 0.0283203 ms
[12/04/2022-02:01:36] [I] GPU Compute Time: min = 30.9525 ms, max = 31.1133 ms, mean = 31.0858 ms, median = 31.0886 ms, percentile(90%) = 31.104 ms, percentile(95%) = 31.1082 ms, percentile(99%) = 31.1133 ms
[12/04/2022-02:01:36] [I] D2H Latency: min = 10.1726 ms, max = 10.1826 ms, mean = 10.1769 ms, median = 10.1768 ms, percentile(90%) = 10.1783 ms, percentile(95%) = 10.1787 ms, percentile(99%) = 10.1826 ms
[12/04/2022-02:01:36] [I] Total Host Walltime: 3.10885 s
[12/04/2022-02:01:36] [I] Total GPU Compute Time: 3.0775 s
[12/04/2022-02:01:36] [I] Explanations of the performance metrics are printed in the verbose logs.

So GPU compute time is about 31ms for a single forward pass, while the reported host latency is about 41ms. The roughly 10ms device-to-host transfer is substantial because this graph returns the full sequence of vocabulary logits. Transfers can sometimes overlap other work, but a dependency on those logits places them on the end-to-end critical path. Thus 31ms is the device-compute measurement and 41ms is the more relevant latency for the exported interface shown here.

Deeper Perf Analysis

Next, using nsight compute and nsight systems I traced and profile the nvidia TRT engines to see what kernels were taking the most time and also to see “overall” how well the engine used the underlying hardware. Combining information from these two sources I looked at Compute and Memory throughput as well as theoretical and actual kernel occupancy on top of an actual kernel trace, producing a trace like below:

Inference of GPT2-Large sequence length=1024, batch=1, on NVIDIA A100

full trace here, you can open this with chrome://tracing in Google Chrome.

Three major observations:

1. Higher sustained compute throughput compared to efficientnet, with fewer bottlenecks.
2. Actual occupancy matches theoretical occupancy much more frequently than efficientnet
3. Very high sustained memory throughput

The trace suggests that high memory bandwidth is an important source of the A100’s latency advantage, though the different TensorRT and OpenVINO kernels prevent a clean attribution.

Since GPUs have a considerable advantage in memory bandwidth over conventional CPUs, I also ran the same exported graph on an Intel machine.

GPT-2 Small Batch inference on Intel Cascade-Lake

For an Intel machine I used the following:

Basic Timing

To get an inference engine optimized for the Intel architecture I used OpenVINO with the following commands

mo --input_model simple_model.onnx 
benchmark_app -m simple_model.xml -hint latency -ip f32 -op f32 -report_type detailed_counters

which resulted in the following (considerably slower) measurements):

Latency:
    Median:     2101.50 ms
    AVG:        2100.02 ms
    MIN:        2085.78 ms
    MAX:        2126.31 ms

The Intel system is much slower for this graph, but the comparison combines hardware, memory system, compiler, and inference-engine differences.

Multiplying peak bandwidth by latency gives a rough implied data-volume scale:

The closer values are consistent with a bandwidth-sensitive workload, though this is not a controlled roofline analysis. CPU platforms with on-package High-Bandwidth Memory (HBM) would make an interesting comparison.

Summary and Conclusions

The results above suggest that this GPT-2 Large full-sequence forward pass at batch size 1 makes meaningful use of the A100, with some room to improve in compute throughput (perhaps by better tuning the shared memory and/or caching in the attention layers).
However since the A100 has excellent memory bandwidth and GPT2 achieves very high sustained memory throughput we see very good performance on the A100, possibly justifying its cost and use (beyond simply the fact that these instances have the most GPU memory and thus fit bigger LLMs).

Intel Cascade Lake performance was poor for this exported graph and software stack. Memory bandwidth is a plausible contributor, and a CPU with on-package HBM would make an interesting follow-up.

Methods and References

Model & checkpoint used and conversion to onnx

I used the Hugging Face gpt2-large configuration, but the script calls from_config and therefore initializes random weights rather than loading the pretrained checkpoint. This preserves the tensor shapes relevant to profiling, but no language-model accuracy is being evaluated.

The rough script for this was

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, onnx
from torchviz import make_dot
import torch
import torch.nn as nn
import onnx
# Configuration
# Load codeparrot tokenizer trained for Python code tokenization
tokenizer = AutoTokenizer.from_pretrained("codeparrot/codeparrot")
# Configure attention scaling/upcasting options supported by this GPT-2 implementation
config_kwargs = {
    "vocab_size": len(tokenizer),
    "scale_attn_by_inverse_layer_idx": True,
    "reorder_and_upcast_attn": True,
}
# Load model config (GPT-2 large in this case)
config = AutoConfig.from_pretrained("gpt2-large", **config_kwargs)
# Initialize new model with config
model = AutoModelForCausalLM.from_config(config)
#Make a fake module to turn off key,value cache so these do not
#get traced during onnx export
class TmpModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.tmp=model
    def forward(self,x):
        return self.tmp(x,use_cache=False)
tmpmodel=TmpModel()
torch.onnx.export(tmpmodel,               # model being run
                  torch.randint(1, len(tokenizer),(1,1024)),  # example input for the model
                  "simple_model.onnx", # where to save the model (can be a file or file-like object)
                  export_params=True,  # store the trained parameter weights inside the model file
                  opset_version=15,    # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names=['input'],   # the model's input names
                  output_names=['output'], # the model's output names
                  training=torch.onnx.TrainingMode.EVAL
                  )

Converting ONNX to TRT Engine

To get fp32 and fp16 engines with extra profiling information like NVTX markers (for use in nsys) I did the following:

trtexec --onnx=simple_model.onnx --saveEngine=gpt2_fp32.engine --profilingVerbosity=detailed --workspace=30000
trtexec --onnx=simple_model.onnx --fp16 --saveEngine=gpt2_fp16.engine --profilingVerbosity=detailed --workspace=30000

Basic timings with trtexec

trtexec --loadEngine=gpt2_fp32.engine --useCudaGraph

Getting nsight systems trace

For the trace I did the following

nsys profile --trace=cuda,cudnn,cublas,osrt,nvtx trtexec --iterations=1 --loadEngine=gpt2_fp32.engine

Getting performance counters with nsight compute

To get performance counters I did

sudo env "PATH=$PATH" ncu --section ComputeWorkloadAnalysis --csv trtexec --loadEngine=gpt2_fp32.engine > cwa.csv
sudo env "PATH=$PATH" ncu --section SpeedOfLight --csv trtexec --loadEngine=gpt2_fp32.engine > sol.csv
sudo env "PATH=$PATH" ncu --section Occupancy --csv trtexec --loadENgine=gpt2_fp32.engine > occupancy.csv

Extracting kernels and start,stop times from nsys systems trace sqlite file

sqlite3 -csv report2.sqlite 'SELECT names.value AS name, start, end FROM CUPTI_ACTIVITY_KIND_KERNEL AS k JOIN StringIds AS names ON k.demangledName = names.id;' > kernels.csv

Combining perf counters and kernels into a single trace file

I combined the perf counter data from nsight compute and the trace of nsight systems to have a visual representation of when the GPU was fully utilized. I did this with a python script you can find here

The resulting file can be opened with chrome://tracing and looks something like below:

References

• permissions setting for nsight compute
• nsight systems tracing
• NVTX tracing
• Nsight compute documentation
• huggingface GPT2 model doc
• some useful LLM definitions
• Huggingface LLM checkpoints
• pytorch onnx support
• chrome tracing file format
• extracting kernels out of nsight systems sqlite
• Openvino documentation
• Openvino benchmarking guide