# Technical Documentation: facebookincubator/AITemplate > ℹ️ **Provenance:** Multi-Tier Ingestion · [CodeWiki Reference](https://codewiki.google/github.com/facebookincubator/AITemplate) ### Overview (`README.md`) # AITemplate [](https://github.com/facebookincubator/AITemplate/blob/main/LICENSE) | [](https://facebookincubator.github.io/AITemplate) | [](https://app.circleci.com/pipelines/github/facebookincubator/AITemplate) [](https://github.com/facebookincubator/AITemplate/actions/workflows/pages.yaml) AITemplate (AIT) is a Python framework that transforms deep neural networks into CUDA (NVIDIA GPU) / HIP (AMD GPU) C++ code for lightning-fast inference serving. AITemplate highlights include: - High performance: close to roofline fp16 TensorCore (NVIDIA GPU) / MatrixCore (AMD GPU) performance on major models, including ResNet, MaskRCNN, BERT, VisionTransformer, Stable Diffusion, etc. - Unified, open, and flexible. Seamless fp16 deep neural network models for NVIDIA GPU or AMD GPU. Fully open source, Lego-style easily extendable high-performance primitives for new model support. Supports a significantly more comprehensive range of fusions than existing solutions for both GPU platforms. ## More about AITemplate ### Excellent Backward Capability AITemplate doesn't depend on third-party libraries or runtimes, such as cuBLAS, cuDNN, rocBLAS, MIOpen, TensorRT, MIGraphX, etc. Each model is compiled into a self-contained portable binary, which can be used on any software environment with the same hardware. ### Horizontal Fusion AITemplate provides unique advanced horizontal fusion. AITemplate can fuse parallel GEMM, LayerNorm, and other operators with different input shapes into a single GPU kernel. ### Vertical Fusion AITemplate provides strong vertical fusion. AITemplate can fuse a large range of operations into TensorCore/MatrixCore operations, such as elementwise operations, reductions, and layout permutations. AITemplate also provides back-to-back style TensorCore / MatrixCore operation fusion. ### Memory Fusion AITemplate provides innovative memory fusions. AITemplate can fuse GEMM, LayerNorm, and other operators, followed by memory operations such as concatenation, split, and slice into a single operator. ### Working w/wo PyTorch The AITemplate-generated Python runtime can take PyTorch tensors as inputs and outputs without an extra copy. For environments without PyTorch, the AITemplate Python/C++ runtime is self-contained. ### Extensions without suffering AITemplate provides a straightforward approach for making an extension in codegen. To add a new operator or a new fused kernel into AITemplate, most of the time one only needs to add two Python files: one for a graph node definition and another for the backend codegen. The CUDA/HIP kernel in a text header file can be directly utilized in the codegen. ## FX2AIT FX2AIT is a Python-based tool that converts PyTorch models into AITemplate (AIT) engine for lightning-fast inference serving. Using FX2AIT's built-in AITLowerer, partial AIT acceleration can be achieved for models with unsupported operators in AITemplate. Key features of FX2AIT include: * Easy Conversion: FX2AIT requires only a PyTorch model and input for conversion, generating an "AITModule" output for inference serving. * Expanded Support: AITemplate does not support all PyTorch operators. FX2AIT's AITLowerer offers a solution for partial AIT conversion for models with unsupported operators. Check the `fx2ait/fx2ait/example/03_lowering_split` for more information. More info can be found from https://github.com/facebookincubator/AITemplate/tree/main/fx2ait. ## Installation **Hardware requirements:** - **NVIDIA**: AIT is only tested on SM80+ GPUs (Ampere etc). Not all kernels work with old SM75/SM70 (T4/V100) GPUs. - **AMD**: AIT is only tested on CDNA2 (MI-210/250) GPUs. There may be compiler issues for old CDNA1 (MI-100) GPUs. ### Clone the code When cloning the code, please use the following command to also clone the submodules: ``` git clone --recursive https://github.com/facebookincubator/AITemplate ``` ### Docker Image We highly recommend using AITemplate with Docker to avoid accidentally using a wrong version of NVCC or HIPCC. - CUDA: `./docker/build.sh cuda` - ROCM: `DOCKER_BUILDKIT=1 ./docker/build.sh rocm` This will build a docker image with tag `ait:latest`. ### From Source The following command will create a Python wheel for AITemplate. Please ensure you have correct CUDA/ROCm compiler installed. - CUDA: CUDA 11.6 - ROCm: We tested on ROCm 5.2.3 with a customized build HIPCC with the command in docker/Dockerfile.rocm#L87-L96 *Incorrect compiler will lead performance regression.* **Please check all submodules are cloned correctly before go to next step.** ``` cd python python setup.py bdist_wheel pip install dist/*.whl --force-reinstall ``` ## Getting Started Check out the [AITemplate Documentation](https://facebookincubator.github.io/AITemplate) for API reference. There are a few tutorials for onboarding: - 01: [How to inference a PyTorch model with AIT](https://facebookincubator.github.io/AITemplate/tutorial/how_to_infer_pt.html) - 02: [How to add an op to AIT codegen](https://facebookincubator.github.io/AITemplate/tutorial/how_to_add_op.html) - 03: [How to visualize AIT's optimization](https://facebookincubator.github.io/AITemplate/tutorial/how_to_visualize.html) ## Examples & Performance AITemplate provides the following model templates & reference performance data on A100/MI-250: - [01_ResNet-50](examples/01_resnet-50/) with PyTorch Image Models (TIMM) - [02_MaskRCNN-FPN](examples/02_detectron2/) with Detectron2 - [03_BERT](examples/03_bert/) with Hugging Face Transformer - [04_Vision Transformer](examples/04_vit/) with PyTorch Image Models (TIMM) - [05_Stable Diffusion](examples/05_stable_diffusion/) with Hugging Face Diffusers ## Release All current development updates can be seen in the AITemplate repository. Releases are not on a set schedule and will only be tagged for significant feature releases. Mid-term plan: - Better dynamic shape support: Focus on the dynamic sequence in Transformers. Add symbolic shape support. - More automatic graph passes: Relief manual rewrite models to obtain the best performance. - Quantization: fp8/int8/int4. - Sparsity pruning for Gemm. - PT2 integration: Aten2AIT is under active development. Long-term plan: - Automatic ONNX, Open-XLA and other format model conversion. - Composable Kernel CPU extension on AVX2/AVX-512 for AMD Epyc CPU. ## Contributing Check our [contributing guide](CONTRIBUTING.md) to learn about how to contribute to the project. ## The Team AITemplate is currently maintained by Meta engineers: [Ying Zhang](https://github.com/ipiszy), [Yang Chen](https://github.com/chenyang78), [Terry Chen](https://github.com/terrychenism), [Mu-Chu Lee](https://github.com/muchulee8), [Max Podkorytov](https://github.com/tenpercent), [Adnan Akhundov](https://github.com/aakhundov). AITemplate is co-created by Meta engineers: [Bing Xu](https://github.com/antinucleon), [Ying Zhang](https://github.com/ipiszy), [Hao Lu](https://github.com/hlu1), [Yang Chen](https://github.com/chenyang78), and [Terry Chen](https://github.com/terrychenism), with major contributions coming from other talented engineers. A non-exhaustive list to mention is Mike Iovine, Mu-Chu Lee, Scott Wolchok, Oleg Khabinov, Shirong Wu, Huamin Li, Hui Guo, Zhijing Li, Max Podkorytov. We also want to thank Andrew Tulloch, Yinghai Lu, Lu Fang for the valuable discussions. FX2AIT and Aten2AIT are co-created and maintained by Meta engineers: [Wei Wei](https://github.com/frank-wei), [Shirong Wu](https://github.com/wushirong) and [Zhijing Li](https://github.com/tissue3). ## Acknowledgements AITemplate team works closely with NVIDIA [CUTLASS](https://github.com/NVIDIA/cutlass) Team (led by Andrew Kerr, Haicheng Wu) and AMD [Composable Kernel](https://github.com/ROCmSoftwarePlatform/composable_kernel) Team (led by Chao Liu, Jing Zhang). We co-designed many advanced GPU optimizations specialized for each platform, and nothing is possible without our close collaboration. ## License AITemplate is licensed under the [Apache 2.0 License](https://github.com/facebookincubator/AITemplate/blob/main/LICENSE). --- ### Overview (`docker/README.md`) # Docker + AITemplate AITemplate provides a Docker image with all test, benchmark, and documentation dependencies installed. ## Build CUDA Docker Image ```bash docker/build.sh cuda``` This will build a CUDA 11 docker image with tag: `ait:latest` ## Build ROCM Docker Image ```DOCKER_BUILDKIT=1 bash docker/build.sh rocm``` This will build a RCOM 5 docker image with tag: `ait:latest` ## Running Unit Tests in Docker ```docker run --gpus all ait:latest bash /AITemplate/tests/nightly/unittest.sh``` ## Launching CUDA Docker ```docker run --gpus all -it ait:latest``` ## Launching ROCM Docker ```docker run -it --network=host --device=/dev/kfd --device=/dev/dri --group-add=video --ipc=host --cap-add=SYS_PTRACE --security-opt seccomp=unconfined ait:latest``` ## Common questions: - Q: When building ROCm Docker, I hit this error ` => ERROR [internal] load metadata for docker.io/library/ubuntu:20.04`, what shall I do? A: Run `docker pull docker.io/library/ubuntu:20.04` to pull base image manually, then re-run `./docker/build.sh rocm` --- ### Overview (`docs/README.md`) # AITemplate Documentation ## Build locally 1. Install AITemplate 2. Install Sphinx ``` pip install autodocsumm pip install sphinx_rtd_theme pip install sphinx_gallery pip install sphinxcontrib-inlinesyntaxhighlight pip install sphinx_toolbox ``` 3. Build HTML ``` make html ``` --- ### Overview (`fx2ait/README.md`) # FX2AIT for AITemplate FX2AIT is a Python-based tool that converts PyTorch models into AITemplate (AIT) engine for lightning-fast inference serving. AITLowerer built on top of FX2AIT is able to perform AIT conversion on PyTorch model with AIT unsupported operators. Model can enjoy partial AIT acceleration using AITLowerer. FX2AIT highlights include: - Automatic Conversion: FX2AIT only need PyTorch model and input as input for conversion. The output can be used for inference serving directly. - Expanded Support: AITemplate doesn't cover all operators PyTorch provides. FX2AIT provided AITLowerer as solution to support partial AIT conversion for models with AIT unsupported operators. For more information, please check example/03_lowering_split. ## Installation **Hardware requirement:** - **NVIDIA**: FX2AIT is based on AIT, thus the hardware requirement is same as AIT. AIT is only tested on SM80+ GPUs (Ampere etc). Not all kernels work with old SM75/SM70 (T4/V100) GPUs. ### From Source The following command will create a Python wheel for AITemplate. Please ensure you have correct CUDA compiler installed. - CUDA: CUDA 11.6 - cuDNN: v8.7.0 for CUDA 11.x download source: https://developer.nvidia.com/rdp/cudnn-download installation guidance: https://docs.nvidia.com/deeplearning/cudnn/install-guide/index.html *Incorrect compiler will lead performance regression.* ``` cd fx2ait python setup.py install ``` ### Docker Image We highly recommend using AITemplate with Docker to avoid accidentally using a wrong version of NVCC or HIPCC. - CUDA: `./docker/build.sh cuda` This will build a docker image with tag `ait:latest`. ## Examples AITemplate provides the following getting started tutorials: - 01: [How to inference a PyTorch Transformer model with FX2AIT](fx2ait/example/01_transformer_model/) - 02: [How to inference a PyTorch vision model with FX2AIT](fx2ait/example/02_vision_model/) - 03: [How to inference a general PyTorch model with AIT unsupported operator using AIT Lowerer](fx2ait/example/03_lowering_split/) ### Run Example and Test Example command: ``` cd fx2ait python example/03_lowering_split/test_lower.py python test/test_ait_lower.py ``` --- ### Overview (`static/README.md`) # About This directory contains all of the C++ sources that are static during AITemplate compilation, including the bulk of the runtime implementation. ## C++ Implementation ### `Model` v.s. `ModelContainer` These are the two main classes involved in the C++ runtime implementation. * The bulk of the runtime implementation is in `Model`. * `ModelContainer` stores a set of shared constants and a collection of `Model`s. Almost all functions in `model_interface.h` forward to a method on `ModelContainer`. When `Run` is invoked, `ModelContainer` looks for an available `Model`, or blocks until one is available (see the section on asynchronous predictions). It then forwards the run request to the runtime. ### Code Structure Some important files: * `include/model_interface.h`: The interface that we expose in the compiled .so * `include/model_container.h`: The bulk of the `ModelContainer` implementation. Some files are generated at compile time. These include: * `model-generated.h`: The implementation for `Model`. * `model_container_base.cu`: A small part of the implementation for `ModelContainer` needs to be codegened. So `ModelContainer` inherits from `ModelContainerBase`, and `ModelContainerBase`'s implementation lives in this file. See `model_container.h` for more details. All codegen templates can be found in `backend/main_templates.py`. The codegen implementation is in `backend/codegen.py`. Note that many of the headers in this directory rely on generated code and thus cannot be `#include`d in external projects. The exception is `model_interface.h`. ## Python `Model` `Model` is a collection of Python bindings to the C++ AIT runtime. This section describes the API. ### `AITData` This class represents a contiguous blob of memory that AIT will use as a tensor. It is simply a named tuple with these fields: * `data_ptr: int`: An **unowned** pointer to **GPU** memory. In general, all of the APIs expect that this pointer will be valid for the entire duration of the call. * `shape: List[int]`: The shape of the tensor. * `dtype: str`: The tensor's dtype; one of `"float32", "float16", "int32", "int64"`. Note that most ops only support float16 at this stage. If using AITemplate with PyTorch, `AITData`s can be constructed with the `torch_to_ait_data` utility: ```python x = torch.randn(3, 3, 3).half().cuda() # Equivalent to AITData(x.data_ptr(), [3, 3, 3], "float16") x_ait = torch_to_ait_data(x) ``` If PyTorch is not available, `Model` provides a set of functions for copying, allocating, and freeing GPU memory. See the docstrings in `compiler/model.py` for more information. ### `run` `run` takes a set of inputs and outputs as `AITData`s. Both arguments can be passed as either an ordered list or a dictionary (mapping name to tensor). ```python # Arguments as a dictionary module.run( {"input0": in0_ait, "input1": in1_ait}, {"output0": out0_ait, "output1": out1_ait}, ) # Arguments as an ordered list. Note that you might need to query # the index mapping. input_name_to_idx = module.get_input_name_to_index_map() output_name_to_idx = module.get_output_name_to_index_map() inputs = [None] * len(input_name_to_idx) outputs = [None] * len(output_name_to_idx) for name in input_name_to_idx: inputs[input_name_to_idx[name]] = ait_inputs[name] for name in output_name_to_idx: outputs[output_name_to_idx[name]] = ait_outputs[name] module.run(inputs, outputs) ``` One important caveat is that the output must be its **maximum** size. This is because of dynamic shapes - the size of the output may vary, but its shape is not inferred until inference time. The maximum shape can be queried with the `get_output_maximum_shape()`: ```python # Can use either name or index. name_to_idx = module.get_output_name_to_idx() max_shape = module.get_output_maximum_shape(name_to_idx["output"]) max_shape = module.get_output_maximum_shape("output") ``` `Model.run` returns a dictionary of output `AITData`s with (possibly dynamic) shapes that the runtime inferred. #### Nullptr Inputs/Outputs In general, inputs are allowed to be null if they are size 0 (e.g. at least one dimension is 0). The runtime enforces this with a check before any kernels are launched. ```cpp If (input_name == nullptr && dim0 * dim1 * … * dimN != 0) { throw std::runtime_error(“input_name cannot be null!”); } ``` This is convenient since torch.data_ptr() returns null for size zero tensors. The dynamic shape computation is skipped if the lower bound of the tensor’s size is positive. #### Constants There are two types of constants in AIT; *bound* and *unbound* constants. A bound constant is known at compile time and may participate in constant folding. Bound constants are copied into GPU memory at model loading time. Values for bound constants may be provided by passing a dictionary (mapping constant name to AIT tensor) to `compile_model`. Unbound constants, on the other hand, do not participate in constant folding and must be provided before running the model. These must be set via `Model.set_constant`: ```python module.set_constant("my_constant", AITData(...)) # The pointer in the the tensor must live for the entire duration of run() module.run(...) ``` Constants are read-only and *shared* with all runtimes in the `ModelContainer`. #### `run_with_tensors` `run_with_tensors` is a convenience method with the same interface as `run`, except it can take lists of `torch.Tensor`s: ```python input0 = torch.randn(input0_shape).cuda().half() output0 = torch.empty(output0_shape).cuda().half() # Returns a dictionary of reshaped outputs result = module.run_with_tensors([input0], [output0]) ``` #### Streams and Asynchronous Predictions A pointer to a stream can optionally be passed to `run`. If none is given, the prediction happens on the default stream 0. If the `sync` argument is set to `True`, the stream is synchronized before `run()` returns. `sync` is `True` by default. Multiple predictions can happen at the same time (on the same or different streams). Under the hood, there is a fixed-size pool of runtime objects. When all the runtimes are used, `run()` blocks until one is available. The size of this pool can be configured with the `num_runtimes` option in `Model`'s constructor. #### CUDA Graph Run also takes a `graph_mode` option. If set to true, the runtime will try to use [CUDA graphs](https://developer.nvidia.com/blog/cuda-graphs/) to run the model. `graph_mode` is not supported on ROCm. The following is a high level overview of how graph mode works: 1) Each `Model` has an internal stream used for graph capturing. The model first runs all ops on this stream in capture mode. No kernel launches happen during this stage. 2) If this is the first run, a graph is instantiated via `cudaGraphInstantiate`. 3) On subsequent runs, we try to avoid the relatively expensive `cudaGraphInstantiate` call by updating the graph executor (`cudaGraphExecUpdate`). However, a new graph may still be instantiated if the topology of the graph somehow changed between runs. 4) Once we have the graph executor, we launch a single kernel on the stream that the user provided to `run()`. Graph mode is mainly beneficial when there are many small kernel launches. A lot of overhead can be avoided since there is only a single kernel launch in graph mode. --- ### Overview (`examples/01_resnet-50/README.md`) # ResNet-50 In this example, we will demo how to use AITemplate for inference on the ResNet-50 model from PyTorch Image Models (TIMM). We will demo two usages: * Using AIT to accelerate PyTorch inference * Using AIT standalone without PyTorch ## Code structure ``` modeling resnet.py # ResNet definition using AIT's frontend API weight_utils.py # Utils to convert TIMM R-50 weights to AIT infer_with_torch.py # Example to accelerate PyTorch, and seamlessly use with other PyTorch code infer_with_numpy.py # Dump TIMM weights to Numpy and use AIT & Numpy without 3rdparties benchmark_pt.py # Benchmark code for PyTorch benchmark_ait.py # Benchmark code for AIT ``` ## Multi-GPU profiling AIT requires to do profiling to decide best algorithms for CUTLASS and CK. To enable multiple GPUs profiling, use the environment variable `CUDA_VISIBLE_DEVICES` on NVIDIA platform and `HIP_VISIBLE_DEVICES` on AMD platform. For example, `CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python3 benchmark_ait.py`. Benchmark is fast once the profilers are built. ## Reference Speed vs PyTorch Eager ### A100-40GB / CUDA 11.6.2 _PT = PyTorch 1.12 Eager_ | Batch size | PT Latency (ms) | PT QPS (im/s) | AIT Latency (ms) | AIT QPS (im/s) | |------------|-----------------|---------------|------------------|----------------| | 1 | 7.68 | 130.29 | 0.58 | 1730.17 | | 2 | 7.16 | 279.36 | 0.62 | 3250.74 | | 4 | 7.17 | 557.68 | 0.69 | 5773.20 | | 8 | 7.02 | 1138.83 | 0.88 | 9104.44 | | 16 | 7.01 | 2280.97 | 1.33 | 12012.81 | | 32 | 7.53 | 4251.30 | 2.40 | 13350.58 | | 64 | 13.98 | 4578.09 | 4.53 | 14140.83 | | 128 | 26.57 | 4816.71 | 8.57 | 14935.82 | | 256 | 50.93 | 5026.40 | 16.58 | 15444.57 | ### MI-250 / ROCm 5.2.3 / HIPCC-10736 _PT = PyTorch 1.12 Eager_ #### 1 GCD | Batch size | PT Latency (ms) | PT QPS (im/s) | AIT Latency (ms) | AIT QPS (im/s) | |------------|-----------------|---------------|------------------|----------------| | 1 | 3.94 | 254.06 | 2.28 | 438.60 | | 2 | 3.89 | 514.48 | 2.25 | 888.89 | | 4 | 3.82 | 1047.11 | 2.38 | 1680.67 | | 8 | 4.40 | 1819.27 | 2.62 | 3053.44 | | 16 | 6.48 | 2468.65 | 3.41 | 4692.08 | | 32 | 10.40 | 3076.97 | 4.86 | 6584.36 | | 64 | 18.35 | 3488.12 | 8.26 | 7748.18 | | 128 | 34.36 | 3724.76 | 15.38 | 8322.50 | | 256 | 65.35 | 3917.29 | 29.62 | 8642.81 | #### 2 GCDs | Batch size | PT Latency (ms) | PT QPS (im/s) | AIT Latency (ms) | AIT QPS (im/s) | |------------|-----------------|---------------|------------------|----------------| | 1 | | | | | | 2 | 3.94 | 507.54 | 2.36 | 848.15 | | 4 | 3.89 | 1028.60 | 2.34 | 1710.94 | | 8 | 3.88 | 2059.41 | 2.70 | 2960.46 | | 16 | 4.56 | 3507.48 | 2.83 | 5663.52 | | 32 | 6.72 | 4762.89 | 3.87 | 8275.98 | | 64 | 10.82 | 5917.63 | 5.26 | 12173.67 | | 128 | 18.79 | 6812.09 | 8.98 | 14247.09 | | 256 | 35.99 | 7112.59 | 16.69 | 15338.58 | ### Note for Performance Results - For NVIDIA A100, our test cluster doesn't allow to lock frequency. We make warm up longer to collect more stable results, but it is expected to have small variance to the results with locked frequency. - To benchmark MI-250, the first step is to run `python3 benchmark_ait.py` to generate all necessary model dynamic library files with single GCD. Then run `./benchmark_mi250.sh {batch_size}` to simulate data parallel execution on 2 GCDs, each GCD is processing half of the batch. - To benchmark MI-250 1 GCD, we lock the frequency with command `rocm-smi -d x --setperfdeterminism 1700`, where `x` is the GPU id. - To benchmark MI-250 2 GCDs, we observed performance regression with rocm perf-determ mode. The 2 GCDs number is running without perf-determ mode set with command `rocm-smi -d x --resetperfdeterminism`, where `x` is the GPU id. - Performance results are what we can reproduce and for reference only. It should not be used for other purposes. --- ### Overview (`examples/02_detectron2/README.md`) # Getting Started with AIT for the Inference of Detectron2 Based Models This document describes the usage of AIT for detectron2 vision models such as mask RCNN and faster RCNN. For an end-to-end example with the API, see `prepare_and_run_rcnn.sh` which covers how to prepare and run inference with `mask_rcnn_R_50_FPN`. ## Create the AIT Model from a Config File 1. Pick a model and its config file from `configs`, for example, `mask_rcnn_R_50_FPN.yaml`. 2. Build the AIT Model by running `compile_model.py` with the config file, for example, ``` cfg=examples/02_detectron2/configs/mask_rcnn_R_50_FPN.yaml CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python3 examples/02_detectron2/compile_model.py \ --config $cfg \ --batch 1 ``` All parameters in the built AIT model are not initialized, and therefore are filled with random values. We will initialize these parameters in the following step (i.e., exporting the weights of the pre-trained model to the AIT model). Check `tmp/mask_rcnn_R_50_FPN/params.json` for the list of parameters in the AIT model and their shapes. ## Download the Detectron2 Pre-trained Model, and Export the Weights to the AIT Model 1. For example, download Detectron2 `mask_rcnn_R_50_FPN` pre-trained model and save it to `tmp/pt_mask_rcnn_R_50_FPN.pkl`: ``` wget https://dl.fbaipublicfiles.com/detectron2/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl -O tmp/pt_mask_rcnn_R_50_FPN.pkl ``` 2. Export the weights from the pre-trained model to AIT model by running `tools/convert_pt2ait.py`: ``` python3 examples/02_detectron2/tools/convert_pt2ait.py \ --d2-weight tmp/pt_mask_rcnn_R_50_FPN.pkl \ --ait-weight tmp/ait_mask_rcnn_R_50_FPN.pt \ --model-name mask_rcnn_R_50_FPN ``` The weights are exported to AIT and saved as `tmp/ait_mask_rcnn_R_50_FPN.pt` for inference run. ## Download Inference DataSet and Run AIT Model 1. For example, download the COCO 2017 Dataset: ``` mkdir -p ~/.torch/datasets/coco wget https://dl.fbaipublicfiles.com/detectron2/annotations/coco/val2017_100.tgz -O ~/.torch/datasets/coco/val2017_100.tgz tar xzf ~/.torch/datasets/coco/val2017_100.tgz -C ~/.torch/datasets/coco && rm -f ~/.torch/datasets/coco/val2017_100.tgz ``` 2. Run inference of the AIT model on the inputs with `demo.py`: ``` python3 examples/02_detectron2/demo.py \ --weight tmp/ait_mask_rcnn_R_50_FPN.pt \ --config examples/02_detectron2/configs/mask_rcnn_R_50_FPN.yaml \ --batch 1 --input "~/.torch/datasets/coco/val2017/*.jpg" \ --confidence-threshold 0.5 \ --display \ --cudagraph ``` ## Multi-GPU profiling AIT requires to do profiling to decide best algorithms for CUTLASS and CK. To enable multiple GPUs profiling, set the environment variable `CUDA_VISIBLE_DEVICES` on NVIDIA platform and `HIP_VISIBLE_DEVICES` on AMD platform with available GPU ids. ## Results _PT = PyTorch 1.12 Eager_ ### A100-40GB / CUDA 11.6 - Input size: 448x608 | Batch size | PT Latency (ms) | PT FPS | AIT Latency (ms) | AIT FPS | |------------|-----------------|--------|------------------|---------| | 1 | 21.70 | 46.09 | 4.40 | 227.27 | | 2 | 29.71 | 67.32 | 6.68 | 299.40 | | 4 | 35.67 | 112.13 | 11.12 | 359.71 | | 8 | 59.71 | 133.98 | 22.24 | 359.71 | | 16 | 112.91 | 141.70 | 36.64 | 436.68 | | 32 | 224.24 | 142.70 | 71.04 | 450.45 | | 64 | 448.84 | 142.59 | 140.16 | 456.62 | - Input size: 800x1344 | Batch size | PT Latency (ms) | PT FPS | AIT Latency (ms) | AIT FPS | |------------|-----------------|--------|------------------|---------| | 1 | 22.99 | 43.50 | 8.50 | 117.65 | | 2 | 34.48 | 58.01 | 13.42 | 149.03 | | 4 | 65.00 | 61.54 | 22.88 | 174.83 | | 8 | 125.25 | 63.87 | 41.44 | 193.05 | | 16 | 246.49 | 64.91 | 78.56 | 203.67 | | 32 | 503.21 | 63.59 | 154.56 | 207.04 | | 64 | OOM | OOM | 304.64 | 210.08 | ### MI-250 / ROCm 5.2.3 / HIPCC-10736 _PT = PyTorch 1.12 Eager_ #### 1 GCDs - Input size: 448x608 | Batch size | PT Latency (ms) | PT FPS | AIT Latency (ms) | AIT FPS | |------------|-----------------|--------|------------------|---------| | 1 | 24.75 | 40.41 | 10.63 | 94.07 | | 2 | 29.28 | 68.30 | 15.96 | 125.31 | | 4 | 42.45 | 94.24 | 26.24 | 152.44 | | 8 | 79.73 | 100.34 | 51.04 | 156.74 | | 16 | 141.84 | 112.81 | 89.12 | 179.53 | | 32 | 284.39 | 112.52 | 161.92 | 197.63 | | 64 | 600.84 | 106.52 | Error | Error | - Input size: 800x1344 | Batch size | PT Latency (ms) | PT FPS | AIT Latency (ms) | AIT FPS | |------------|-----------------|--------|------------------|---------| | 1 | 26.80 | 37.31 | 19.23 | 52.00 | | 2 | 43.61 | 45.86 | 30.28 | 66.05 | | 4 | 98.88 | 40.45 | 51.56 | 77.58 | | 8 | 189.45 | 42.23 | 98.80 | 80.97 | | 16 | 389.94 | 41.03 | 177.28 | 90.25 | | 32 | 807.22 | 39.64 | 333.44 | 95.97 | | 64 | 1768.66 | 36.19 | Error | Error | #### 2 GCDs - Input size: 448x608 | Batch size | AIT Latency (ms) | AIT FPS | |------------|------------------|---------| | 1 | | | | 2 | 12.78 | 156.49 | | 4 | 20.66 | 193.61 | | 8 | 32.16 | 248.76 | | 16 | 61.52 | 260.08 | | 32 | 106.08 | 301.66 | | 64 | 194.24 | 329.49 | - Input size: 800x1344 | Batch size | AIT Latency (ms) | AIT FPS | |------------|------------------|---------| | 1 | | | | 2 | 22 | 90.91 | | 4 | 34 | 117.65 | | 8 | 55.52 | 144.09 | | 16 | 104.48 | 153.14 | | 32 | 190.24 | 168.21 | | 64 | 362.88 | 176.37 | ### Sample outputs ### Note for Performance Results - For NVIDIA A100, our test cluster doesn't allow to lock frequency. We make warm up longer to collect more stable results, but it is expected to have small variance to the results with locked frequency. - To benchmark MI-250, the first step is to run `python3 benchmark_ait.py` to generate all necessary model dynamic library files with single GCD. Then run `./benchmark_mi250.sh {batch_size}` to simulate data parallel execution on 2 GCDs, each GCD is processing half of the batch. - To benchmark MI-250 1 GCD, we lock the frequency with command `rocm-smi -d x --setperfdeterminism 1700`, where `x` is the GPU id. - To benchmark MI-250 2 GCDs, we observed performance regression with rocm perf-determ mode. The 2 GCDs number is running without perf-determ mode set with command `rocm-smi -d x --resetperfdeterminism`, where `x` is the GPU id. - Performance results are what we can reproduced. It should not be used for other purposes. --- ### Overview (`examples/03_bert/README.md`) # BERT This directory contains an AIT demo for the [BERT language representation model](https://huggingface.co/docs/transformers/v4.22.1/en/model_doc/bert). Only `bert-base-uncased` is included. ## Prerequisites Install the dependencies: ``` python3 -m pip install transformers click torch ``` ## Benchmarking To run a basic benchmark, use `benchmark.py`: ``` python3 examples/03_bert/benchmark_ait.py ``` There are two options for hidden activations, `gelu` and `fast_gelu` (`fast_gelu` by default). `gelu` is not supported on AMD hardware yet. ``` python3 examples/03_bert/benchmark_ait.py --activation gelu python3 examples/03_bert/benchmark_ait.py --activation fast_gelu ``` The batch size and sequence length can also be configured via the command line: ``` python3 examples/03_bert/benchmark_ait.py --batch-size 1 --seq-length 128 ``` PyTorch eager mode benchmarks can also be run: ``` python3 examples/03_bert/benchmark_pt.py ``` To benchmark BERT embeddings, run benchmark with `--encoders-only False` ## Quick Demo To run a quick demo with a simple prompt, use `demo.py`: ``` python3 examples/03_bert/demo.py --prompt "The quick brown fox jumps over the lazy dog." ``` The demo prints out the resulting logits. The demo only works with sequence length <= 512. ## Multi-GPU profiling AIT requires to do profiling to decide best algorithms for CUTLASS and CK. To enable multiple GPUs profiling, use the environment variable `CUDA_VISIBLE_DEVICES` on NVIDIA platform and `HIP_VISIBLE_DEVICES` on AMD platform. ## Reference Speed vs PyTorch Eager _PT = PyTorch 1.12 Eager_ _OOM = Out of Memory_ ### A100-40GB / CUDA 11.6.2 - Sequence length 64 | Batch size | PT Latency (ms) | PT QPS (seq/s) | AIT Latency (ms) | AIT QPS (seq/s) | |------------|-----------------|----------------|------------------|-----------------| | 1 | 7.96 | 125.65 | 0.71 | 1399.64 | | 2 | 8.38 | 238.59 | 0.74 | 2719.15 | | 4 | 8.29 | 482.30 | 0.80 | 4994.37 | | 8 | 8.51 | 939.97 | 0.95 | 8439.67 | | 16 | 8.09 | 1978.47 | 1.41 | 11385.85 | | 32 | 9.19 | 3481.34 | 2.23 | 14357.58 | | 64 | 9.12 | 7016.80 | 4.14 | 15458.15 | | 128 | 14.52 | 8814.57 | 8.00 | 15991.44 | | 256 | 27.75 | 9224.39 | 15.99 | 16006.79 | - Sequence length 128 | Batch size | PT Latency (ms) | PT QPS (seq/s) | AIT Latency (ms) | AIT QPS (seq/s) | |------------|-----------------|----------------|------------------|-----------------| | 1 | 8.02 | 124.72 | 0.78 | 1281.52 | | 2 | 8.29 | 241.22 | 0.85 | 2364.94 | | 4 | 8.51 | 470.29 | 0.99 | 4044.33 | | 8 | 8.12 | 985.72 | 1.43 | 5600.93 | | 16 | 9.22 | 1735.20 | 2.21 | 7232.47 | | 32 | 9.11 | 3512.80 | 4.17 | 7677.82 | | 64 | 15.29 | 4184.93 | 8.05 | 7949.06 | | 128 | 29.44 | 4347.33 | 16.03 | 7987.11 | | 256 | 56.34 | 4543.88 | 31.57 | 8109.08 | - Sequence length 384 | Batch size | PT Latency (ms) | PT QPS (seq/s) | AIT Latency (ms) | AIT QPS (seq/s) | |------------|-----------------|----------------|------------------|-----------------| | 1 | 8.72 | 114.73 | 1.63 | 611.91 | | 2 | 8.31 | 240.73 | 1.97 | 1013.19 | | 4 | 8.64 | 463.10 | 2.55 | 1569.23 | | 8 | 9.32 | 858.70 | 3.95 | 2025.62 | | 16 | 13.90 | 1151.03 | 6.80 | 2354.21 | | 32 | 26.72 | 1197.74 | 13.30 | 2405.46 | | 64 | 51.02 | 1254.34 | 26.68 | 2398.95 | | 128 | 100.26 | 1276.67 | 51.60 | 2480.67 | | 256 | OOM | OOM | 101.55 | 2520.81 | - Sequence length 1024 | Batch size | PT Latency (ms) | PT QPS (seq/s) | AIT Latency (ms) | AIT QPS (seq/s) | |------------|-----------------|----------------|------------------|-----------------| | 1 | 9.74 | 102.65 | 2.20 | 454.12 | | 2 | 11.38 | 175.75 | 4.15 | 481.95 | | 4 | 13.61 | 293.90 | 8.36 | 478.44 | | 8 | 25.79 | 310.15 | 12.53 | 638.53 | | 16 | 49.91 | 320.59 | 21.61 | 740.48 | | 32 | 97.00 | 329.91 | 42.84 | 746.88 | | 64 | 191.14 | 334.83 | 83.95 | 762.39 | | 128 | OOM | OOM | 163.96 | 780.70 | | 256 | OOM | OOM | 324.22 | 789.58 | - Sequence length 4096 | Batch size | PT Latency (ms) | PT QPS (seq/s) | AIT Latency (ms) | AIT QPS (seq/s) | |------------|-----------------|----------------|------------------|-----------------| | 1 | 32.82 | 30.47 | 18.23 | 54.87 | | 2 | 65.25 | 30.65 | 35.64 | 56.11 | | 4 | 128.73 | 31.07 | 103.67 | 38.58 | | 8 | OOM | OOM | 119.45 | 66.98 | | 16 | OOM | OOM | 166.25 | 96.24 | | 32 | OOM | OOM | 333.98 | 95.81 | | 64 | OOM | OOM | 662.29 | 96.63 | | 128 | OOM | OOM | 1313.77 | 97.43 | | 256 | | | | | ### MI-250 / ROCm 5.2.3 / HIPCC-10736 #### 1 GCD - Sequence length 64 | Batch size | PT Latency (ms) | PT QPS (seq/s) | AIT Latency (ms) | AIT QPS (seq/s) | |------------|-----------------|----------------|------------------|-----------------| | 1 | 5.72 | 174.72 | 2.78 | 359.88 | | 2 | 5.96 | 335.38 | 2.87 | 697.76 | | 4 | 5.85 | 684.16 | 2.85 | 1404.31 | | 8 | 6.15 | 1300.72 | 3.15 | 2540.72 | | 16 | 6.14 | 2605.40 | 3.78 | 4231.12 | | 32 | 7.73 | 4138.06 | 5.34 | 5993.50 | | 64 | 14.38 | 4451.07 | 9.10 | 7030.42 | | 128 | 26.18 | 4889.95 | 16.45 | 7780.40 | | 256 | 49.95 | 5125.04 | 31.90 | 8023.98 | - Sequence length 128 | Batch size | PT Latency (ms) | PT QPS (seq/s) | AIT Latency (ms) | AIT QPS (seq/s) | |------------|-----------------|----------------|------------------|-----------------| | 1 | 5.76 | 173.55 | 2.68 | 373.03 | | 2 | 6.06 | 330.18 | 2.87 | 697.33 | | 4 | 5.96 | 670.65 | 3.02 | 1324.91 | | 8 | 6.03 | 1326.23 | 3.65 | 2194.62 | | 16 | 9.35 | 1711.55 | 4.98 | 3212.12 | | 32 | 16.46 | 1943.61 | 8.48 | 3775.22 | | 64 | 30.83 | 2075.74 | 15.44 | 4146.40 | | 128 | 58.74 | 2179.24 | 30.57 | 4187.68 | | 256 | 115.27 | 2220.87 | 59.28 | 4318.61 | - Sequence length 384 | Batch size | PT Latency (ms) | PT QPS (seq/s) | AIT Latency (ms) | AIT QPS (seq/s) | |------------|-----------------|----------------|------------------|-----------------| | 1 | 5.78 | 172.87 | 2.97 | 336.14 | | 2 | 6.02 | 332.30 | 3.45 | 579.89 | | 4 | 8.00 | 499.85 | 4.68 | 854.16 | | 8 | 13.79 | 580.01 | 7.47 | 1070.24 | | 16 | 24.39 | 656.06 | 13.04 | 1226.77 | | 32 | 45.56 | 702.33 | 24.26 | 1318.80 | | 64 | 87.84 | 728.57 | 47.87 | 1336.92 | | 128 | 172.57 | 741.71 | 95.22 | 1344.26 | | 256 | 352.27 | 726.71 | 185.94 | 1376.78 | - Sequence length 1024 | Batch size | PT Latency (ms) | PT QPS (seq/s) | AIT Latency (ms) | AIT QPS (seq/s) | |------------|-----------------|----------------|------------------|-----------------| | 1 | 6.86 | 145.71 | 4.20 | 237.84 | | 2 | 12.41 | 161.21 | 5.82 | 343.62 | | 4 | 22.25 | 179.80 | 10.20 | 392.26 | | 8 | 41.94 | 190.73 | 18.91 | 423.05 | | 16 | 81.03 | 197.45 | 37.86 | 422.60 | | 32 | 159.06 | 201.19 | 71.65 | 446.62 | | 64 | 321.51 | 199.06 | 148.86 | 429.95 | | 128 | OOM | OOM | 277.53 | 461.21 | | 256 | OOM | OOM | 563.07 | 454.65 | - Sequence length 4096 | Batch size | PT Latency (ms) | PT QPS (seq/s) | AIT Latency (ms) | AIT QPS (seq/s) | |------------|-----------------|----------------|------------------|-----------------| | 1 | 49.89 | 20.04 | 16.18 | 61.81 | | 2 | 93.22 | 21.45 | 30.67 | 65.21 | | 4 | 183.57 | 21.79 | 66.78 | 59.90 | | 8 | 366.57 | 21.82 | 117.49 | 68.09 | | 16 | OOM | OOM | 231.15 | 69.22 | | 32 | OOM | OOM | 459.46 | 69.65 | | 64 | OOM | OOM | 1031.86 | 62.02 | | 128 | | | | | | 256 | | | | | #### 2 GCDs - Sequence length 64 | Batch size | PT Latency (ms) | PT QPS (seq/s) | AIT Latency (ms) | AIT QPS (seq/s) | |------------|-----------------|----------------|------------------|-----------------| | 1 | | | | | | 2 | 5.52 | 362.55 | 2.80 | 714.99 | | 4 | 6.04 | 661.73 | 2.89 | 1385.05 | | 8 | 6.07 | 1317.20 | 2.82 | 2835.38 | | 16 | 6.02 | 2659.82 | 3.29 | 4866.99 | | 32 | 6.09 | 5257.45 | 3.83 | 8352.10 | | 64 | 8.53 | 7506.95 | 5.81 | 11013.02 | | 128 | 15.34 | 8346.14 | 10.00 | 12806.23 | | 256 | 28.44 | 9002.30 | 18.92 | 13528.13 | - Sequence length 128 | Batch size | PT Latency (ms) | PT QPS (seq/s) | AIT Latency (ms) | AIT QPS (seq/s) | |------------|-----------------|----------------|------------------|-----------------| | 1 | | | | | | 2 | 5.58 | 358.62 | 2.68 | 745.20 | | 4 | 6.20 | 644.91 | 2.83 | 1411.55 | | 8 | 6.08 | 1316.09 | 3.21 | 2492.88 | | 16 | 5.89 | 2716.79 | 3.86 | 4144.50 | | 32 | 9.86 | 3247.03 | 5.41 | 5915.33 | | 64 | 17.71 | 3614.25 | 9.64 | 6640.53 | | 128 | 32.74 | 3909.15 | 17.81 | 7186.25 | | 256 | 62.73 | 4080.77 | 35.73 | 7165.20 | - Sequence length 384 | Batch size | PT Latency (ms) | PT QPS (seq/s) | AIT Latency (ms) | AIT QPS (seq/s) | |------------|-----------------|----------------|------------------|-----------------| | 1 | | | | | | 2 | 5.57 | 358.88 | 3.09 | 647.71 | | 4 | 6.12 | 653.83 | 3.62 | 1104.69 | | 8 | 8.35 | 958.19 | 4.94 | 1620.06 | | 16 | 14.29 | 1119.38 | 8.29 | 1930.01 | | 32 | 26.10 | 1226.17 | 14.96 | 2139.07 | | 64 | 50.01 | 1279.72 | 28.22 | 2268.02 | | 128 | 97.55 | 1312.15 | 55.94 | 2288.37 | | 256 | 193.00 | 1326.44 | 111.27 | 2300.68 | - Sequence length 1024 | Batch size | PT Latency (ms) | PT QPS (seq/s) | AIT Latency (ms) | AIT QPS (seq/s) | |------------|-----------------|----------------|------------------|-----------------| | 1 | | | | | | 2 | 6.80 | 294.16 | 4.36 | 458.93 | | 4 | 13.01 | 307.55 | 6.43 | 622.23 | | 8 | 23.39 | 341.99 | 11.52 | 694.52 | | 16 | 44.45 | 359.94 | 21.83 | 732.90 | | 32 | 87.23 | 366.84 | 43.73 | 731.77 | | 64 | 172.92 | 370.12 | 82.92 | 771.85 | | 128 | 352.09 | 363.54 | 173.14 | 739.29 | | 256 | OOM | OOM | 322.97 | 792.64 | - Sequence length 4096 | Batch size | PT Latency (ms) | PT QPS (seq/s) | AIT Latency (ms) | AIT QPS (seq/s) | |------------|-----------------|----------------|------------------|-----------------| | 1 | | | | | | 2 | 54.67 | 36.58 | 18.31 | 109.23 | | 4 | 104.19 | 38.39 | 35.09 | 113.99 | | 8 | 206.62 | 38.72 | 77.03 | 103.86 | | 16 | 412.58 | 38.78 | 133.59 | 119.77 | | 32 | OOM | OOM | 263.40 | 121.49 | | 64 | OOM | OOM | 524.11 | 122.11 | | 128 | OOM | OOM | 1186.20 | 107.91 | | 256 | | | | | ### Note Performance Results - For NVIDIA A100, our test cluster doesn't allow to lock frequency. We make warm up longer to collect more stable results, but it is expected to have small variance to the results with locked frequency. - To benchmark MI-250, the first step is to run `python3 benchmark_ait.py` to generate all necessary model dynamic library files with single GCD. Then run `./benchmark_mi250.sh {batch_size}` to simulate data parallel execution on 2 GCDs, each GCD is processing half of the batch. - To benchmark MI-250 1 GCD, we lock the frequency with command `rocm-smi -d x --setperfdeterminism 1700`, where `x` is the GPU id. - To benchmark MI-250 2 GCDs, we observed performance regression with rocm perf-determ mode. The 2 GCDs number is running without perf-determ mode set with command `rocm-smi -d x --resetperfdeterminism`, where `x` is the GPU id. - PyTorch Eager result doesn't reflect [BetterTransformer](https://pytorch.org/blog/a-better-transformer-for-fast-transformer-encoder-inference/), mainly due to BetterTransformer integration to TIMM/Transformer package is not yet landed. - Performance results are what we can reproduced. It should not be used for other purposes. --- ### Overview (`examples/04_vit/README.md`) # Vision Transformer (VIT) In this example, we will demo how to lower a pretrained Vision Transformer from TIMM, and run inference in AITemplate. We tested on two variants of Vision Transformer: Base version with 224x224 input / patch 16, and Large version with 384x384 input / patch 16. ## Code structure ``` modeling vision_transformer.py # VIT definition using AIT's frontend API weight_utils.py # Utils to convert TIMM VIT weights to AIT verification.py # Numerical verification between TIMM and AIT benchmark_pt.py # Benchmark code for PyTorch benchmark_ait.py # Benchmark code for AITemplate ``` ## Reference Speed vs PyTorch Eager ### A100-40GB / CUDA 11.6.2 _PT = PyTorch 1.12 Eager_ - vit_base_patch16_224 | Batch size | PT Latency (ms) | PT QPS (im/s) | AIT Latency (ms) | AIT QPS (im/s) | |------------|-----------------|---------------|------------------|----------------| | 1 | 4.95 | 202.15 | 1.02 | 979.31 | | 2 | 5.26 | 380.43 | 1.15 | 1735.64 | | 4 | 5.51 | 726.08 | 1.57 | 2543.72 | | 8 | 5.56 | 1439.03 | 2.20 | 3642.16 | | 16 | 8.59 | 1863.35 | 3.64 | 4396.74 | | 32 | 15.95 | 2006.62 | 6.51 | 4916.93 | | 64 | 31.48 | 2032.77 | 12.67 | 5052.52 | | 128 | 59.86 | 2138.35 | 25.10 | 5099.77 | | 256 | 115.00 | 2226.10 | 48.55 | 5273.03 | - vit_large_patch16_384 | Batch size | PT Latency (ms) | PT QPS (im/s) | AIT Latency (ms) | AIT QPS (im/s) | |------------|-----------------|---------------|------------------|----------------| | 1 | 9.88 | 101.17 | 3.84 | 260.21 | | 2 | 11.90 | 168.02 | 5.87 | 340.98 | | 4 | 21.20 | 188.66 | 11.49 | 348.09 | | 8 | 39.33 | 203.43 | 19.09 | 419.07 | | 16 | 76.00 | 210.54 | 36.19 | 442.08 | | 32 | 147.24 | 217.33 | 70.03 | 456.93 | | 64 | 291.00 | 219.93 | 135.25 | 473.21 | | 128 | 578.99 | 221.08 | 267.09 | 479.24 | | 256 | 1204.16 | 212.60 | 538.97 | 474.98 | ### MI-250 / ROCm 5.2.3 / HIPCC-10736 _PT = PyTorch 1.12 Eager_ #### 1 GCD - vit_base_patch16_224 | Batch size | PT Latency (ms) | PT QPS (im/s) | AIT Latency (ms) | AIT QPS (im/s) | |------------|-----------------|---------------|------------------|----------------| | 1 | 3.54 | 282.12 | 3.49 | 286.26 | | 2 | 4.43 | 451.73 | 3.78 | 528.84 | | 4 | 6.09 | 657.02 | 4.05 | 986.95 | | 8 | 9.65 | 829.27 | 5.31 | 1507.06 | | 16 | 16.62 | 962.98 | 8.50 | 1882.72 | | 32 | 29.87 | 1071.25 | 14.43 | 2218.07 | | 64 | 56.58 | 1131.08 | 26.52 | 2413.45 | | 128 | 110.28 | 1160.73 | 51.62 | 2479.69 | | 256 | 217.07 | 1179.35 | 102.82 | 2489.89 | - vit_large_patch16_384 | Batch size | PT Latency (ms) | PT QPS (im/s) | AIT Latency (ms) | AIT QPS (im/s) | |------------|-----------------|---------------|------------------|----------------| | 1 | 12.90 | 77.51 | 9.70 | 103.05 | | 2 | 22.42 | 89.19 | 13.40 | 149.29 | | 4 | 38.16 | 104.83 | 22.12 | 180.86 | | 8 | 70.58 | 113.35 | 38.46 | 208.00 | | 16 | 136.28 | 117.40 | 70.44 | 227.15 | | 32 | 261.97 | 122.15 | 138.14 | 231.65 | | 64 | 541.90 | 118.10 | 270.01 | 237.02 | | 128 | 1108.36 | 115.49 | 534.97 | 239.27 | | 256 | 2213.09 | 115.68 | 1063.24 | 240.77 | #### 2 GCDs - vit_base_patch16_224 | Batch size | PT Latency (ms) | PT QPS (im/s) | AIT Latency (ms) | AIT QPS (im/s) | |------------|-----------------|---------------|------------------|----------------| | 1 | | | | | | 2 | 3.49 | 572.95 | 3.59 | 556.55 | | 4 | 4.11 | 974.26 | 3.97 | 1006.80 | | 8 | 5.88 | 1359.64 | 4.23 | 1889.44 | | 16 | 9.75 | 1641.06 | 5.71 | 2800.69 | | 32 | 17.55 | 1823.03 | 9.34 | 3426.32 | | 64 | 31.31 | 2043.79 | 16.24 | 3940.53 | | 128 | 60.33 | 2121.64 | 30.97 | 4133.14 | | 256 | 117.96 | 2170.29 | 59.82 | 4279.21 | - vit_large_patch16_384 | Batch size | PT Latency (ms) | PT QPS (im/s) | AIT Latency (ms) | AIT QPS (im/s) | |------------|-----------------|---------------|------------------|----------------| | 1 | | | | | | 2 | 12.73 | 157.07 | 10.52 | 190.13 | | 4 | 22.97 | 174.12 | 14.94 | 267.82 | | 8 | 39.78 | 201.08 | 24.55 | 325.85 | | 16 | 74.95 | 213.48 | 43.95 | 364.07 | | 32 | 146.18 | 218.91 | 82.04 | 390.06 | | 64 | 283.04 | 226.12 | 162.62 | 393.55 | | 128 | 583.03 | 219.54 | 313.34 | 408.51 | | 256 | 1197.56 | 213.77 | 621.71 | 411.77 | ### Note for Performance Results - For NVIDIA A100, our test cluster doesn't allow to lock frequency. We make warm up longer to collect more stable results, but it is expected to have small variance to the results with locked frequency. - To benchmark MI-250, the first step is to run `python3 benchmark_ait.py` to generate all necessary model dynamic library files with single GCD. Then run `./benchmark_mi250.sh {batch_size}` to simulate data parallel execution on 2 GCDs, each GCD is processing half of the batch. - To benchmark MI-250 1 GCD, we lock the frequency with command `rocm-smi -d x --setperfdeterminism 1700`, where `x` is the GPU id. - To benchmark MI-250 2 GCDs, we observed performance regression with rocm perf-determ mode. The 2 GCDs number is running without perf-determ mode set with command `rocm-smi -d x --resetperfdeterminism`, where `x` is the GPU id. - PyTorch Eager result doesn't reflect [BetterTransformer](https://pytorch.org/blog/a-better-transformer-for-fast-transformer-encoder-inference/), mainly due to BetterTransformer integration to TIMM/Transformer package is not yet landed. - Performance results are what we can reproduce. It should not be used for other purposes. --- ### Overview (`examples/05_stable_diffusion/README.md`) ## Stable Diffusion Example In this example, we show how to build fast AIT modules for CLIP, UNet, VAE models, and benchmark/run them. ### Build Dependencies First, clone, build, and install AITemplate [per the README instructions](https://github.com/facebookincubator/AITemplate#clone-the-code). This AIT stable diffusion example depends on `diffusers`, `transformers`, `torch` and `click`. You could install them using `pip`. Verify the library versions. We have tested transformers==4.25, diffusers==0.11[torch] and torch==1.12. ``` >>> import transformers >>> transformers.__version__ '4.25.0' >>> import diffusers >>> diffusers.__version__ '0.11.0' >>> import torch >>> torch.__version__ '1.12.0+cu113' ``` ### Download the diffusers pipeline files Optionally, you can use Hugging Face access token. You can register in Hugging Face Hub to obtain an access token for the Stable Diffusion weights. See [user access tokens](https://huggingface.co/docs/hub/security-tokens) for more info. Your access tokens are listed in your [Hugging Face account settings](https://huggingface.co/settings/tokens). stable-diffusion model has two variants - base and regular. For example: - `stabilityai/stable-diffusion-2-1-base` - image resolution 512x512 - `stabilityai/stable-diffusion-2-1` - image resolution 768x768 ``` python3 scripts/download_pipeline.py \ --model-name "stabilityai/stable-diffusion-2-1-base" # Optionally, you can use access token --token ACCESS_TOKEN ``` ### Build AIT modules for CLIP, UNet, VAE Build the AIT modules by running `compile.py`. Set correct width and height depending on the model variant ``` python3 scripts/compile.py --width 512 --height 512 ``` It generates three folders: `./tmp/CLIPTextModel`, `./tmp/UNet2DConditionModel`, `./tmp/AutoencoderKL`. In each folder, there is a `test.so` file which is the generated AIT module for the model. #### Alternative build script ``` python3 scripts/compile_alt.py --width 64 1536 --height 64 1536 --batch-size 1 4 --clip-chunks 6 ``` This compiles modules with dynamic shape. In the example, modules will work with width in range 64-1536px, batch sizes 1-4. Clip chunks refers to the number of tokens accepted by UNet in multiples of 77, 1 chunk = 77 tokens, 3 chunks = 231 tokens. By default, `compile_alt.py` does not include model weights (constants) with the compiled module, to include the model weights in the compiled module use `--include-consants True`. #### Alternative pipeline The original pipeline requires a diffusers model local dir, and relies directly on `StableDiffusionPipeline`. This pipeline builds similar functionality without directly using `StableDiffusionPipeline`, and is capable of loading model weights from either diffusers or [CompVis](https://huggingface.co/CompVis) models to compiled aitemplate modules. * AITemplate modules are created * Model weights are loaded, converted/mapped, then applied to AITemplate module * Tokenizer is created from `openai/clip-vit-large-patch14`. * Scheduler is created from `hf-hub-or-path`. * Loading CLIPTextModel from `ckpt` requires the appropriate `hf-hub-or-path` to be specified i.e. `runwayml/stable-diffusion-v1-5` for SD1.x checkpoints, `stabilityai/stable-diffusion-2-1` for SD2.x checkpoints. ``` python3 scripts/demo_alt.py \ --hf-hub-or-path runwayml/stable-diffusion-v1-5 \ --ckpt v1-5-pruned-emaonly.ckpt python3 scripts/demo_alt.py \ --hf-hub-or-path stabilityai/stable-diffusion-2-1 \ --ckpt v2-1_768-ema-pruned.ckpt ``` `--ckpt` takes preference over `--hf-hub-or-path` if both are specified #### Multi-GPU profiling AIT needs to do profiling to select the best algorithms for CUTLASS and CK. To enable multiple GPUs for profiling, use the environment variable `CUDA_VISIBLE_DEVICES` on NVIDIA platform and `HIP_VISIBLE_DEVICES` on AMD platform. ### Benchmark This step is optional. You can run `benchmark.py` to measure throughput for each of the subnets. Benchmark script supports base model variant only for now - 512x512 ``` python3 src/benchmark.py ``` ### Verify This step is optional. You can verify numerical correctness for each of the subnets. ``` HUGGINGFACE_AUTH_TOKEN=ACCESS_TOKEN python3 -m unittest src/test_correctness.py ``` ### Run Models Run AIT models with an example image: Set correct width and height depending on the model variant ``` python3 scripts/demo.py --width 512 --height 512 ``` Img2img demo: Set correct width and height depending on the model variant ``` python3 scripts/demo_img2img.py --width 512 --height 512 ``` Check the resulted image: `example_ait.png` ### Sample outputs Command: `python3 scripts/demo.py --prompt "Mountain Rainier in van Gogh's world"` Command: `python3 scripts/demo.py --prompt "Sitting in a tea house in Japan with Mount Fuji in the background, sunset professional portrait, Nikon 85mm f/1.4G"` Command: `scripts/demo.py --prompt "A lot of wild flowers with North Cascade Mountain in background, sunset professional photo, Unreal Engine"` ## Results _PT = PyTorch 1.12 Eager_ _OOM = Out of Memory_ ### A100-40GB / CUDA 11.6, 50 steps | Module | PT Latency (ms) | AIT Latency (ms) | |----------|-----------------|------------------| | CLIP | 9.48 | 0.87 | | UNet | 60.52 | 22.47 | | VAE | 47.78 | 37.43 | | Pipeline | 3058.27 | 1282.98 | - PT: 17.50 it/s - AIT: 42.45 it/s ### RTX 3080-10GB / CUDA 11.6, 50 steps | Module | PT Latency (ms) | AIT Latency (ms) | |----------|-----------------|------------------| | CLIP | OOM | 0.85 | | UNet | OOM | 40.22 | | VAE | OOM | 44.12 | | Pipeline | OOM | 2163.43 | - PT: OOM - AIT: 24.51 it/s ### MI-250 1 GCD, 50 steps | Module | PT Latency (ms) | AIT Latency (ms) | |----------|-----------------|------------------| | CLIP | 6.16 | 2.98 | | UNet | 78.42 | 62.18 | | VAE | 63.83 | 164.50 | | Pipeline | 4300.16 | 3476.07 | - PT: 12.43 it/s - AIT: 15.60 it/s ## Batched Version ### A100-40GB / CUDA 11.6 - Stable Diffusion with AIT batch inference, 50 steps | Batch size | PT Latency (ms) | AIT Latency (ms) | |--------------|------------------|------------------| | 1 | 3058.27 | 1282.98 | | 3 | 7334.46 | 3121.88 | | 8 | 17944.60 | 7492.81 | | 16 | OOM | 14931.95 | - AIT Faster rendering, 25 steps | Batch size | AIT Latency (ms) | AVG im/s | |------------|------------------|----------| | 1 | 695 | 0.69 | | 3 | 1651 | 0.55 | | 8 | 3975 | 0.50 | | 16 | 7906 | 0.49 | ## IMG2IMG ### A100-40GB / CUDA 11.6, 40 steps | Module | PT Latency (ms) | AIT Latency (ms) | |----------|-----------------|------------------| | Pipeline | 4163.60 | 1785.46 | ### Note for Performance Results - For all benchmarks we render the images of size 512x512 - For img2img model we only support fix input 512x768 by default, stay tuned for dynamic shape support - For NVIDIA A100, our test cluster doesn't allow to lock frequency. We make warm up longer to collect more stable results, but it is expected to have small variance to the results with locked frequency. - To benchmark MI-250 1 GCD, we lock the frequency with command `rocm-smi -d x --setperfdeterminism 1700`, where `x` is the GPU id. - Performance results are what we can reproduced & take reference. It should not be used for other purposes. --- ### Overview (`examples/08_esrgan/README.md`) ## ESRGAN Example In this example, we show how to build fast AIT modules for ESRGAN models, and benchmark/run them. ### Build Dependencies First, clone, build, and install AITemplate [per the README instructions](https://github.com/facebookincubator/AITemplate#clone-the-code). This AIT ESRGAN example depends on `torch`, `click` and optionally `safetensors`. You could install them using `pip`. ### Download the ESRGAN model We have tested the following ESRGAN models. [RealESRGAN_x4plus](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth) Model architecture: ``` num_in_ch: 3, num_out_ch: 3, num_feat: 64, num_block: 23, num_grow_ch: 32, scale: 4, ``` [RealESRGAN_x4plus_anime_6B](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth) Model architecture: ``` num_in_ch: 3, num_out_ch: 3, num_feat: 64, num_block: 6, num_grow_ch: 32, scale: 4, ``` [RealESRGAN_x2plus](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth) Model architecture: ``` num_in_ch: 3, num_out_ch: 3, num_feat: 64, num_block: 23, num_grow_ch: 32, scale: 2, ``` A database of ESRGAN models can be found [here](https://upscale.wiki/wiki/Model_Database). Safetensors versions are supported. ### Build AIT modules for ESRGAN Build the AIT modules by running `compile.py`. ``` Usage: compile.py [OPTIONS] Options: --model-path TEXT model path. supports torch or safetensors --width ... Minimum and maximum width --height ... Minimum and maximum height --batch-size ... Minimum and maximum batch size --include-constants BOOLEAN include constants (model weights) with compiled model --num-in-ch INTEGER Number of in channels --num-out-ch INTEGER Number of out channels --num-feat INTEGER Number of intermediate features --num-block INTEGER Number of RRDB layers --num-grow-ch INTEGER Number of channels for each growth --scale INTEGER Scale --use-fp16-acc BOOLEAN use fp16 accumulation --convert-conv-to-gemm BOOLEAN convert 1x1 conv to gemm --work-dir TEXT Work directory --model-name TEXT Model name --help Show this message and exit. ``` Use `--num-in-ch`, `--num-out-ch`, `--num-feat`, `--num-block`, `--num-grow-ch`, `--scale` options to set the ESRGAN model architecture. The default values are for `RealESRGAN_x4plus` architecture. `--width` and `--height` require a minimum and maximum value, the compiled module supports the range of resolutions. However, with 2x model architecture, only static shape is supported, the maximum value for each dimension is used. Defaults are `64` and `1024`. `--batch-size` is supported for 4x model architecture, provide minimum and maximum values. Default is `1 1`. Use `--include-constants False` to compile the module without model weights. AIT modules are compatible with all ESRGAN models with the same model architecture. This can simplify deployment by compiling a module without model weights then applying weights at runtime by using AIT mapped weights (see `map_rrdb` in `./modeling/rrdbnet.py`) with the module's `set_many_constants_with_tensors`. In our tests an ESRGAN module compiled with weights is approximately `~38MB`, and `~6.5MB` without. Examples: ``` python compile.py --model-path "RealESRGAN_x4plus.safetensors" ``` ``` python compile.py --model-path "RealESRGAN_x4plus_anime_6B.pth" --num-block 6 --model-name RealESRGAN_x4plus_anime_6B ``` ``` python compile.py --model-path "RealESRGAN_x2plus.pth" --scale 2 --model-name RealESRGAN_x2plus --width 512 512 --height 512 512 ``` #### Multi-GPU profiling AIT needs to do profiling to select the best algorithms for CUTLASS and CK. To enable multiple GPUs for profiling, use the environment variable `CUDA_VISIBLE_DEVICES` on NVIDIA platform and `HIP_VISIBLE_DEVICES` on AMD platform. ### Run Models `demo.py` provides example usage of ESRGAN modules. ``` Usage: demo.py [OPTIONS] Options: --module-path TEXT the AIT module path --input-image-path TEXT path to input image --output-image-path TEXT path to output image --scale INTEGER Scale of ESRGAN model --help Show this message and exit. ``` `--scale` must match the scale of the model architecture. Limitations: * Demo does not support multiple images/batch size. * Demo does not support images with alpha channel. --- ### Overview (`tests/ci_profile_cache/README.md`) # Profiling Database for CI (Deprecated) Profile Cache DB for CI is deprecated. Now CI will select the algorithm with the smallest tiling size and smallest alignments for CI. The selection function is defined at: `backend/target.py: Target:select_minimal_algo` and specialized in each backend target implementation. --- ### Overview (`fx2ait/fx2ait/example/01_transformer_model/README.md`) # Transfomer encoder In this example, we will demo how to use FX2AIT for inference on the transformer encoder block from pytorch. ## Code structure ``` test_transformer_encoder.py # Transformer encoder block definition using torch API ../benchmark_utils.py # Accuracy verification and Benchmark code for FX2AIT ``` ## How to Use FX2AIT allows users to directly define a torch model, while fx2ait converter does the conversion for the usage. Therefore the encoder can be defined normally as ``` class EncoderBlock(torch.nn.Module): def __init__(self, input_dim, num_heads, dim_feedforward, dropout=0.0): super().__init__() # Attention layer self.attn = torch.nn.MultiheadAttention( embed_dim=input_dim, num_heads=num_heads, batch_first=True, ) # # Two-layer MLP self.linear_net = torch.nn.Sequential( torch.nn.Linear(input_dim, dim_feedforward), torch.nn.Dropout(dropout), torch.nn.ReLU(inplace=True), torch.nn.Linear(dim_feedforward, input_dim), ) # Layers to apply in between the main layers self.norm1 = torch.nn.LayerNorm(input_dim) self.norm2 = torch.nn.LayerNorm(input_dim) self.dropout = torch.nn.Dropout(dropout) def forward(self, x): # Attention part attn_out, _ = self.attn(query=x, key=x, value=x) # return attn_out x = x + self.dropout(attn_out) x = self.norm1(x) # MLP part linear_out = self.linear_net(x) x = x + self.dropout(linear_out) x = self.norm2(x) return x ``` To run the test and benchmark, ``` python fx2ait/fx2ait/example/01_transformer_model/test_transformer_encoder.py ``` ## Reference Speed vs PyTorch Eager ### A100-40GB / CUDA 11.6.2 _PT = PyTorch 1.12 Eager_ | Batch size | PT Latency (s) | PT QPS (im/s) | AIT Latency (ms) | AIT QPS (im/s) | Speedup | |------------|----------------|---------------|------------------|----------------|------------| | 1 | 0.00043845 | 2280.75893 | 0.0001806 | 5537.09872 | 2.42774396 | | 8 | 0.00047376 | 8443.01343 | 0.0002221 | 18009.9959 | 2.13312416 | | 16 | 0.00085377 | 18740.4255 | 0.00050193 | 31876.7364 | 1.90096119 | | 32 | 0.00150154 | 21311.3919 | 0.00069578 | 45991.3908 | 2.15806602 | | 64 | 0.00296888 | 21556.9773 | 0.00138113 | 46338.7065 | 2.14959202 | | 128 | 0.00530519 | 24127.3232 | 0.00261813 | 48889.8245 | 2.02632609 | | 256 | 0.01015745 | 25203.1791 | 0.00516545 | 49560.0242 | 1.96641955 | | 512 | 0.02023099 | 25307.7086 | 0.01034528 | 49491.1828 | 1.95557739 | ### Note for Performance Results - For NVIDIA A100, our test cluster doesn't allow us to lock frequency. We make warm up longer to collect more stable results, but it is expected to have small variance to the results with locked frequency. - Performance results are what we can reproduce and for reference only. It should not be used for other purposes. --- ### Overview (`fx2ait/fx2ait/example/02_vision_model/README.md`) # ResNet-18 In this example, we will demo how to use FX2AIT for inference on the ResNet-18 model from torchvision. ## Code structure ``` test_vision_model.py # ResNet definition using torch API ../benchmark_utils.py # Accuracy verification and Benchmark code for FX2AIT ``` ## How to Use FX2AIT allows users to directly define a torch model, while fx2ait converter does the conversion for the usage. Therefore the definition of model is as simple as ``` class TestModule(torch.nn.Module): def __init__(self): super().__init__() self.mod = torchvision.models.resnet18() def forward(self, x): return self.mod(x) ``` Notice that because AIT supports channel last, while pytorch supports channel first operation, FX2AIT automatically performs this layout conversion for you. To run the test and benchmark, ``` python fx2ait/fx2ait/example/02_vision_model/test_vision_model.py ``` ## Reference Speed vs PyTorch Eager ### A100-40GB / CUDA 11.6.2 _PT = PyTorch 1.12 Eager_ | Batch size | PT Latency (s) | PT QPS (im/s) | AIT Latency (ms) | AIT QPS (im/s) | Speedup | |------------|----------------|---------------|------------------|----------------|------------| | 1 | 0.00349264 | 286.316562 | 0.00052888 | 1890.78465 | 6.60382564 | | 8 | 0.00382057 | 2093.93053 | 0.0007766 | 10301.2714 | 4.91958606 | | 16 | 0.00351062 | 4557.59936 | 0.00098235 | 16287.4093 | 3.57368167 | | 32 | 0.00321071 | 9966.64244 | 0.00166504 | 19218.8053 | 1.92831291 | | 256 | 0.01670636 | 15323.5057 | 0.01181243 | 21672.0808 | 1.41430305 | | 512 | 0.03276252 | 15627.6137 | 0.02347752 | 21808.0915 | 1.39548442 | ### Note for Performance Results - For NVIDIA A100, our test cluster doesn't allow us to lock frequency. We make warm up longer to collect more stable results, but it is expected to have small variance to the results with locked frequency. - Performance results are what we can reproduce and for reference only. It should not be used for other purposes. --- ### Overview (`fx2ait/fx2ait/example/03_lowering_split/README.md`) # AIT Lowerer Now let's go back the example of encoder. Imagine the cases that 1) You want to test the effect of a particular op, say MultiheadAttention on the entire module. 2) There is some special op that AIT doesn't support. AIT actually provide an **automatic** Lowerer to split the graph into subgraphs and run interpreter, so that AIT only run the part it can handle and leave other to AITemplate. In this example, we will demo how to use AitLowerer for inference on any models. ## Code structure ``` test_lowerr.py # Splited transformer encoder block to illustrate the usage of AitLowerer. ../benchmark_utils.py # Accuracy verification and Benchmark code for FX2AIT ../lower/ lower.py # Lower interface, which integrates lowering passes of Split subgraph and AIT Interpreter ait_splitter.py # Splitter to split graph into submodules ait_setting.py # Lowering settings ``` ## How to Use To skip an operation can be extremely easy. One just need to register in the method function `@torch.fx.wrap` ``` @torch.fx.wrap def unsupported_attention_op(f, x): attn_out, _ = f(x, x, x) return attn_out ``` Then at forward stage, call the function. ``` class LowerModule(torch.nn.Module): def __init__(self, input_dim, num_heads, dim_feedforward, dropout=0.0): super().__init__() self.attn = torch.nn.MultiheadAttention( embed_dim=input_dim, num_heads=num_heads, batch_first=True, ) ... def forward(self, x): # Unsupported op will not be lowered to AIT backend. attn_out = unsupported_attention_op(self.attn, x) ... ``` Then AIT won't deal with that part. ``` lowerer = AitLowerer.create( LowerSettings( workdir="/tmp", name="test_ait_lower", min_acc_module_size=0, ) ) lowered = lowerer(model, inputs) lower_output = lowered(*inputs) ``` The mechanism is that Acc tracer allows user to register wrap function so that Acc won't deal with it. Then our splitter will split the them into subgraph: _run_on_gpu_0 for pytorch Eager mode and _run_on_acc_1 for AIT, where _run_on_gpu_0 contains torch.nn.MultiheadAttention and _run_on_acc_1 contains the rest of the model. Finally, interpreter will be called for the AIT subgraph. (_run_on_acc_1) *Notice that our splitter only split subgraphs with more than 10 ops, since otherwise the subgraph is too small.* To run the test and benchmark, ``` python fx2ait/fx2ait/example/03_lowering_split/test_lower.py ``` ## Reference Speed vs PyTorch Eager ### A100-40GB / CUDA 11.6.2 _PT = PyTorch 1.12 Eager_ | Batch size | PT Latency (s) | PT QPS (im/s) | AIT Latency (ms) | AIT QPS (im/s) | Speedup | |------------|----------------|---------------|------------------|----------------|------------| | 1 | 0.00065761 | 1520.66428 | 0.00076476 | 1307.59981 | 0.85988724 | | 4 | 0.00090687 | 4410.77681 | 0.00079056 | 5059.68597 | 1.14711902 | | 16 | 0.00249116 | 6422.69897 | 0.00200686 | 7972.66574 | 1.24132639 | | 32 | 0.00473638 | 6756.209 | 0.00396992 | 8060.62008 | 1.19306849 | | 64 | 0.00914742 | 6996.51201 | 0.00754977 | 8477.07749 | 1.2116148 | | 128 | 0.0178672 | 7163.96537 | 0.01501702 | 8523.66305 | 1.1897968 | | 256 | 0.03554306 | 7202.53192 | 0.02998132 | 8538.65123 | 1.18550689 | | 512 | 0.07118476 | 7192.55069 | 0.06006168 | 8524.56943 | 1.18519421 | From the example, we learn without AIT's multihead attention module, the speedup will be degraded to 1.2x compared to Pytorch Eager. ### Note for Performance Results - For NVIDIA A100, our test cluster doesn't allow us to lock frequency. We make warm up longer to collect more stable results, but it is expected to have small variance to the results with locked frequency. - Performance results are what we can reproduce and for reference only. It should not be used for other purposes. --- ### Source Debughints (`docs/source/debughints.rst`) Debug Hints =========== AITemplate is a new project under active development. We have a rich test set to avoid bugs but don't be surprised if there is anything unexpected. Here are some helpful tips we learned during the development of AITemplate: 1. Once the codegen for op which requires profiling is changed, remember to delete old profilers (usually located at workdir), and flush the cache by either deleting `~/.aitemplate` or setting the environment variable `FLUSH_PROFILE_CACHE=1`. 2. Check the pseudo code/visualization generated by each optimization pass if some optimization behaves in unexpected way. 3. Always do the numerical test, from small to large, to make sure the entire model is correct. 4. Try to make the new fusion subgraph work in a manual way, then try to add an automatic pass to rewrite the graph with the fused subgraph. --- ### Source Index (`docs/source/index.rst`) AITemplate Documentation ======================== AITemplate (AIT) is a Python framework that transforms deep neural networks into CUDA (NVIDIA GPU) / HIP (AMD GPU) C++ code for lightning-fast inference serving. AITemplate highlights include: * High performance: close to roofline fp16 TensorCore (NVIDIA GPU) / MatrixCore (AMD GPU) performance on major models, including ResNet, MaskRCNN, BERT, VisionTransformer, Stable Diffusion, etc. * Unified, open, and flexible. Seamless fp16 deep neural network models for NVIDIA GPU or AMD GPU. Fully open source, Lego-style easily extendable high-performance primitives for new model support. Supports a significantly more comprehensive range of fusions than existing solutions for both GPU platforms. .. toctree:: :maxdepth: 1 :caption: Getting Started install/index .. toctree:: :maxdepth: 1 :caption: User Guide tutorial/index debughints .. toctree:: :maxdepth: 1 :caption: Runtime Design runtime/index .. toctree:: :maxdepth: 1 :caption: Architecture Guide arch/index .. toctree:: :maxdepth: 1 :caption: Reference Guide reference/index reference/env genindex --- ### Source Arch Index (`docs/source/arch/index.rst`) Design and Architecture ======================= .. toctree:: :maxdepth: 1 philosophy Stay tuned for more... --- ### Source Arch Philosophy (`docs/source/arch/philosophy.rst`) Design Philosophy ================== KISS (Keep it simple and stupid) -------------------------------- AITemplate avoids deep IR lowering stacks to reduce the system's complexity. A highly modularized, multiple backend codegen system written in pure Python directly attacks the pain point in high-performance GPU inference. Pragmatism ---------- AITemplate provides a PyTorch-style frontend to enable engineers to manually match the PyTorch model & weights to AITemplate for optimization. Using it is less painful than debugging different lowering IR stacks, especially for complex models such as MaskRCNN. We believe most of the neural network workload can be decoupled. For example, most of the network can be decoupled into Encoder, Decoder, and Decoder logics. For encoder and decoder, it is a computation-bounded problem. For decoder logic, it may involve more control flows. By using divide and conquer, we left the decoder logic part to C++ or Python rather than build a unified language / IR stack as a silver bullet. --- ### Source Install Index (`docs/source/install/index.rst`) Installing AITemplate ===================== Using Docker ------------ The easiest way to get started is to use Docker. Using docker is able to avoid performance regression caused by incorrect version of NVCC and HIPCC. To use docker, we provide a bash script to build the docker image. - CUDA: .. code-block:: bash ./docker/build.sh cuda - ROCM: .. code-block:: bash DOCKER_BUILDKIT=1 ./docker/build.sh rocm This will build a docker image with tag `ait:latest`. To launch the docker container - CUDA: .. code-block:: bash docker run --gpus all -it ait:latest - ROCM: .. code-block:: bash docker run -it --network=host --device=/dev/kfd --device=/dev/dri --group-add=video --ipc=host --cap-add=SYS_PTRACE --security-opt seccomp=unconfined ait:latest AITemplate will be installed as a Python package in Python 3.8. There will be also a copy of the source code and examples at `/AITemplate`. Installing as a Standard Python Package --------------------------------------- Before installing AITemplate, first make sure you have correct hardware and software environment. - Hardware - NVIDIA: AIT is only tested on SM80+ GPUs (Ampere etc). - AMD: AIT is only tested on CDNA2 (MI-210/250) GPUs. .. warning:: - Not all kernels work with old SM75/SM70 (T4/V100) GPUs. - There may be compiler issues for old CDNA1 (MI-100) GPUs. - Software - NVIDIA: CUDA 11.6 - AMD: ROCm 5.2, with HIPCC 10736 (commit `b0f4678b9058a4ae00200dfb1de0da5f2ea84dcb`) .. warning:: - Incorrect compiler version may lead to performance regression. - Instruction for building HIPCC 10736 can be founded in `docker/Dockerfile.rocm`. When cloning the code, please use the following command to clone the submodules: .. code-block:: bash git clone --recursive https://github.com/facebookincubator/AITemplate .. warning:: Please check that all submodules are cloned correctly before the next step. Then build the Python wheel package and install it: .. code-block:: bash cd python python setup.py bdist_wheel pip install dist/aitemplate-*.whl --- ### Source Reference Backend (`docs/source/reference/backend.rst`) aitemplate.backend =========================== aitemplate.backend.task_runner ------------------------------ .. automodule:: aitemplate.backend.task_runner :members: :imported-members: :exclude-members: OrderedDict :autosummary: aitemplate.backend.builder -------------------------- .. automodule:: aitemplate.backend.builder :members: :imported-members: :exclude-members: BaseRunner, Target :autosummary: aitemplate.backend.codegen --------------------------- .. automodule:: aitemplate.backend.codegen :members: :imported-members: :exclude-members: Tensor, Target :autosummary: aitemplate.backend.profiler_cache ---------------------------------- .. automodule:: aitemplate.backend.profiler_cache :members: :imported-members: :exclude-members: :autosummary: aitemplate.backend.profiler_runner ----------------------------------- .. automodule:: aitemplate.backend.profiler_runner :members: :imported-members: :exclude-members: Target, Task, namedtuple, BaseRunner :autosummary: aitemplate.backend.registry ---------------------------- .. automodule:: aitemplate.backend.registry :members: :imported-members: :exclude-members: :autosummary: aitemplate.backend.target -------------------------- .. automodule:: aitemplate.backend.target :members: :imported-members: :exclude-members: :autosummary: --- ### Source Reference Compiler (`docs/source/reference/compiler.rst`) aitemplate.compiler ============================== base ------------------------ .. automodule:: aitemplate.compiler.base :members: :imported-members: :exclude-members: ABC, Enum, abstructmethod, dataclass, pformat, reduce :autosummary: tensor_accessor ----------------------------------- .. automodule:: aitemplate.compiler.tensor_accessor :members: :imported-members: :exclude-members: IntImm, IntVar, Tensor, pformat :autosummary: compiler ---------------------------- .. automodule:: aitemplate.compiler.compiler :members: :imported-members: :exclude-members: IntImm, IntVar, Tensor, pformat, DynamicProfileStrategy :autosummary: model ---------------------------- .. automodule:: aitemplate.compiler.model :members: :imported-members: :exclude-members: NamedTuple, TypeVar :autosummary: --- ### Source Reference Cuda (`docs/source/reference/cuda.rst`) aitemplate.backend.cuda =========================== target_def ---------- .. automodule:: aitemplate.backend.cuda.target_def :members: :imported-members: :exclude-members: Path, ProfileCacheDB, TargetType :autosummary: --- ### Source Reference Env (`docs/source/reference/env.rst`) Environment Variables ===================== AITemplate uses environment variables to control the behavior of codegen and profiling. The environment variables used in AITemplate are listed here. Codegen ------- **NUM_BUILDERS**: The number of CPU jobs running in parallel during codegen. It controls both the profiler codegen and the final .so codegen. It's set to 12 in NIGHTLY jobs. Internally, it's set to 12 for normal tests and 24 for heavy tests. By default, the builder uses all the available CPUs for building. **AIT_RECOMPILE**: If set to "0", it skips compilation for the .so and reuses the previously compiled ones. It is used to speed up local testing. The default value is "1" to always recompile. **AIT_NDEBUG**: If set to "1", compile with `NDEBUG`, disabling debug assertions. Recommended for production builds. "1" by default. **AIT_COMPILER_OPT**: The optimization level for a compiler, which is directly passed to the host compiler command line. AITemplate host code may be very light in certain cases, so there is nothing to optimize for a host compiler. Thus, there is no need to make host compiler perform time costly optimizations. It may be very useful to use "-O0" value for debugging GPU kernels. "-O3" by default. **AIT_NVCC_CCBIN**: nvcc host compiler (ccbin). **AIT_ENABLE_CUDA_LTO**: If set to "1", nvcc will use LTO flags during compilation. Default value is "0". **AIT_TIME_COMPILATION**: If set to "1", time each make command at the compilation time. This helps us to do compilation time analysis. Requires to install `time `_ package. **AIT_MULTISTREAM_MODE**: Controls multi-stream mode. Default mode is "0". * If set to "0", then no multistreaming is used. * If set to "1", then a simple multistreaming is used (iteratively track a wavefront of independent operators and execute ones). **AIT_MULTISTREAM_EXTRA_STREAMS**: Specifies the number of additional streams used. Default value is "4". **AIT_MULTISTREAM_MAX_MEM_PARALLEL_OPS**: Maximum number of parallel operators used in memory planning for simple multi-stream mode. Default value is "99999999" (basically, unlimited). **AIT_USE_CMAKE_COMPILATION**: (An experimental feature) If set to "1", then `cmake` will used instead of `make`. This allows to build AITemplate using MSVC Compiler + MSBuild on Windows, and it works for linux as well. This builder does not support many features (such as caching) yet. But it allows to generate a cmake project that can be loaded to a modern IDE. Default value is "0". **AIT_ENABLE_STANDALONE**: Enable standalone test and benchmark executable generation. Default value is "0" (disabled). If set to "1", this will generate a "test" executable that may be used to run standalone tests and benchmarks. This standalone executable is also well suited for running through debuggers and/or profiling tools, as it does not pull in python and pytorch as dependencies, unlike most python unit tests. **AIT_ENABLE_PTXAS_INFO**: Set this to "1" to enable the generation and logging of verbose ( tuning-relevant ) information about CUDA ptx assembly code produced by the CUDA compiler nvcc. Intermediate ptx files, annotated with C++ source info will be written to the build directory. In addition, this flag enables warnings about CUDA register spilling and resource usage. **AIT_CUDA_DEBUG_LEVEL**: Configure level of CUDA debug information. Defaults to no debug info. This may either be a string with options passed to nvcc ( for example "-g -G" or "-lineinfo" ) or a CUDA debug level from "0" (default, no debug info), "1" ( "-lineinfo" ) include source code line information. Ideal for profiling with ncu/nsight-compute, "2" full debug information (**warning**: this disables all optimizations, regardless of other settings) **AIT_ENABLE_CUDA_SOURCE_NAVIGATION_FIX**: (Only supported by FBCUDA target so far): When this flag is enabled by setting it to "1" (it is disabled by default), every *.cu file in build dirs into a corresponding *.cu.h file and create a *.cu file which just includes this file. This fixes code navigation issues in some IDE's which don't treat .cu files as C++ files and disable code navigation. **AIT_ENABLE_INCLUDE_FROM_SOURCETREE**: (Only supported by FBCUDA target so far) When this flag is enabled by setting it to "1" (it is disabled by default), the target will create an in-place build which tries to directly reference the include paths within the AITemplate source tree. This helps to iterate faster during native Kernel/Operator development and debugging. Profiling --------- **CACHE_DIR**: The directory for the profiling cache. If unset, it defaults to `~/.aitemplate`. **FLUSH_PROFILE_CACHE**: If set to "1", it removes the cache file and recreates an empty one. **DISABLE_PROFILER_CODEGEN**: Normally in CI we randomly choose two profilers to codegen. If set to "1", this flag disables profiler codegen completely to speed up long running tests so that the tests don't time out. The default value is "0". **CUDA_VISIBLE_DEVICES**: This one is from CUDA itself. It's used to set the number of GPU devices available for profiling. Set to "0,1,2,3,4,5,6,7" to speed up profiling. For benchmarking, it's useful to set to a particular device to lower noise. **HIP_VISIBLE_DEVICES**: This one is from ROCm itself. It's used to set the number of GPU devices available for profiling. Set to "0,1,2,3,4,5,6,7" to speed up profiling. For benchmarking, it's useful to set to a particular device to lower noise. **FORCE_PROFILE**: If set to "1", it will do profiling regardless in_ci_env and disable_profiler_codegen. For non-NIGHTLY CI, we do not do profiling, and we could use FORCE_PROFILE=1 in these CI to do runs with codegen, compile, and profile. **COMBINE_PROFILER_MULTI_SOURCES**: Whether to combine multiple profiler sources per target. "0" - Disabled, "1" - Enabled (default). **FORCE_ONE_PROFILER_SOURCE_PER_TARGET**: Whether to combine multiple profiler sources per target into one. "0" - Disabled (default), "1" - Enabled. OSS CI ------ **CI_FLAG**: It is set to "CIRCLECI" in OSS CI to indicate we're in OSS CI environment. The behavior of the profiler and codegen is different in CI to speed up testing. Profiling itself for gemm/conv ops is disabled in CI. But we still compile two random profilers to make sure the profiler codegen is not broken. **AIT_BUILD_DOCS**: If set to "1", it will create a fake CUDA target to enable doc building in Github Actions. Miscellaneous ------------- **LOGLEVEL**: It is used to control the logging level in Python. The default value is "INFO". "DEBUG" is useful for debugging. **AIT_PLOT_SHORTEN_TENSOR_NAMES**: If set to "1", shorten too long tensor names for a plot of a model graph, thus making a plot much easier to analyze visually. "0" by default. **AIT_USE_FAST_MATH**: If set to "0", no fast math option will be used for the device code generation. Default value is "1". **AIT_USE_TANH_FOR_SIGMOID**: If set to "1", tanh will be used to approximate sigmoid during device code generation. Default value is "0". --- ### Source Reference Frontend (`docs/source/reference/frontend.rst`) aitemplate.frontend ==================== .. automodule:: aitemplate.frontend.nn :members: :imported-members: :exclude-members: :autosummary: .. automodule:: aitemplate.frontend.tensor :members: :imported-members: :exclude-members: :autosummary: --- ### Source Reference Index (`docs/source/reference/index.rst`) Python API ========== .. toctree:: :maxdepth: 2 compiler ops transform backend cuda rocm frontend testing utils --- ### Source Reference Ops (`docs/source/reference/ops.rst`) aitemplate.compiler.ops ======================== .. automodule:: aitemplate.compiler.ops :members: :imported-members: :exclude-members: Tensor, TensorAccessor, Enum, Operator, IntImm, IntVar, IntVarTensor, wrap_dim :autosummary: --- ### Source Reference Rocm (`docs/source/reference/rocm.rst`) aitemplate.backend.rocm =========================== target_def ---------- .. automodule:: aitemplate.backend.rocm.target_def :members: :imported-members: :exclude-members: :autosummary: --- ### Source Reference Testing (`docs/source/reference/testing.rst`) aitemplate.testing ================== detect_target ------------- .. automodule:: aitemplate.testing.detect_target :members: :imported-members: :exclude-members: CUDA, ROCM, Popen :autosummary: benchmark_pt ------------ .. automodule:: aitemplate.testing.benchmark_pt :members: :imported-members: :exclude-members: CUDA, ROCM, Popen :autosummary: benchmark_ait ------------- .. automodule:: aitemplate.testing.benchmark_ait :members: :imported-members: :exclude-members: CUDA, ROCM, Popen :autosummary: --- ### Source Reference Transform (`docs/source/reference/transform.rst`) aitemplate.compiler.transform ============================== apply_padding ------------------------------------------- .. automodule:: aitemplate.compiler.transform.apply_padding :members: :imported-members: :exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm :autosummary: bind_constants ------------------------------------------- .. automodule:: aitemplate.compiler.transform.bind_constants :members: :imported-members: :exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm :autosummary: constant_folding ------------------------------------------- .. automodule:: aitemplate.compiler.transform.constant_folding :members: :imported-members: :exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm, AITData, replace_tensor :autosummary: fuse_conv_elementwise ------------------------------------------- .. automodule:: aitemplate.compiler.transform.fuse_conv_elementwise :members: :imported-members: :exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm, AITData, replace_tensor :autosummary: fuse_group_ops ------------------------------------------- .. automodule:: aitemplate.compiler.transform.fuse_group_ops :members: :imported-members: :exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm, AITData, replace_tensor, all_static_dimensions :autosummary: fuse_mm_elementwise ------------------------------------------- .. automodule:: aitemplate.compiler.transform.fuse_mm_elementwise :members: :imported-members: :exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm, AITData, replace_tensor, FuncEnum, elementwise, gemm_rcr, gemm_rcr_bias, gemm_rcr_bias_swish, copy_tensor_attributes, extract_only_one_op, get_patterns, remove_single_tensor_op_from_sorted_graph, sanitize_sorted_graph :autosummary: fuse_ops ------------------------------------------- .. automodule:: aitemplate.compiler.transform.fuse_ops :members: :imported-members: :exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm, AITData, replace_tensor, FuncEnum, elementwise, gemm_rcr, gemm_rcr_bias, gemm_rcr_bias_swish, copy_tensor_attributes, extract_only_one_op, get_patterns, remove_single_tensor_op_from_sorted_graph, sanitize_sorted_graph, layernorm_sigmoid_mul :autosummary: fuse_parallel_gemms ------------------------------------------- .. automodule:: aitemplate.compiler.transform.fuse_parallel_gemms :members: :imported-members: :exclude-members: TensorAccessor, is_static_dimension, DimInfo, IntImm, Operator, Source, Tensor, gemm, AITData, replace_tensor, FuncEnum, elementwise, gemm_rcr, gemm_rcr_bias, gemm_rcr_bias_swish, copy_tensor_attributes, extract_only_one_op, get_patterns, remove_single_tensor_op_from_sorted_graph, sanitize_sorted_graph, layernorm_sigmoid_mul :autosummary: fuse_permute_bmm ------------------------------------------- .. automodule:: aitemplate.compiler.transform.fuse_permute_bmm :members: :imported-members: :exclude-members: copy_src_op_attributes, remove_tensor_from_sorted_graph, bmm_ccr, bmm_crr, bmm_rcr, bmm_rrr, gemm_rrr, gemm_rrr_bias, permute021, TensorAccessor, is_static_dimension, DimInfo, IntImm, Operator, Source, Tensor, gemm, AITData, replace_tensor, FuncEnum, elementwise, gemm_rcr, gemm_rcr_bias, gemm_rcr_bias_swish, copy_tensor_attributes, extract_only_one_op, get_patterns, remove_single_tensor_op_from_sorted_graph, sanitize_sorted_graph, layernorm_sigmoid_mul :autosummary: fuse_split ------------------------------------------- .. automodule:: aitemplate.compiler.transform.fuse_split :members: :imported-members: _fuse_split_and_strided_op :exclude-members: IntVar, copy_src_op_attributes, remove_tensor_from_sorted_graph, bmm_ccr, bmm_crr, bmm_rcr, bmm_rrr, gemm_rrr, gemm_rrr_bias, permute021, TensorAccessor, is_static_dimension, DimInfo, IntImm, Operator, Source, Tensor, gemm, AITData, replace_tensor, FuncEnum, elementwise, gemm_rcr, gemm_rcr_bias, gemm_rcr_bias_swish, copy_tensor_attributes, extract_only_one_op, get_patterns, remove_single_tensor_op_from_sorted_graph, sanitize_sorted_graph, layernorm_sigmoid_mul :autosummary: mark_param_tensor ------------------------------------------- .. automodule:: aitemplate.compiler.transform.mark_param_tensor :members: :imported-members: :exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm :autosummary: memory_planning ------------------------------------------- .. automodule:: aitemplate.compiler.transform.memory_planning :members: :imported-members: :exclude-members: TensorUsageRecord, Workspace, assign_offsets_to_views_and_outputs, greedy_by_size_memory_planning, DimInfo, IntImm, Operator, Source, Tensor, gemm, defaultdict, dataclass :autosummary: name_graph ------------------------------------------- .. automodule:: aitemplate.compiler.transform.name_graph :members: :imported-members: :exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm :autosummary: optimize_graph ------------------------------------------- .. automodule:: aitemplate.compiler.transform.optimize_graph :members: :imported-members: :exclude-members: transform_strided_ops, transform_special_ops, transform_odd_alignment, transform_memory_ops, fuse_permute_bmm, fuse_parallel_gemms, fuse_mm_elementwise, apply_padding, fuse_conv_elementwise, fuse_group_ops, DimInfo, IntImm, Operator, Source, Tensor, gemm :autosummary: profile ------------------------------------------- .. automodule:: aitemplate.compiler.transform.profile :members: :imported-members: :exclude-members: DynamicProfileStrategy, DimInfo, IntImm, Operator, Source, Tensor, gemm :autosummary: refine_graph ------------------------------------------- .. automodule:: aitemplate.compiler.transform.refine_graph :members: :imported-members: :exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm :autosummary: remove_no_ops ------------------------------------------- .. automodule:: aitemplate.compiler.transform.remove_no_ops :members: :imported-members: :exclude-members: IntVar, is_singleton_dimension, DimInfo, IntImm, Operator, Source, Tensor, gemm :autosummary: remove_unused_ops ------------------------------------------- .. automodule:: aitemplate.compiler.transform.remove_unused_ops :members: :imported-members: :exclude-members: deque, DimInfo, IntImm, Operator, Source, Tensor, gemm :autosummary: toposort ------------------------------------------- .. automodule:: aitemplate.compiler.transform.toposort :members: :imported-members: :exclude-members: DimInfo, IntImm, Operator, Source, Tensor, gemm :autosummary: transform_memory_ops ------------------------------------------- .. automodule:: aitemplate.compiler.transform.transform_memory_ops :members: :imported-members: :exclude-members: TensorAccessor, DimInfo, IntImm, Operator, Source, Tensor, gemm :autosummary: transform_odd_alignment ------------------------------------------- .. automodule:: aitemplate.compiler.transform.transform_odd_alignment :members: :imported-members: :exclude-members: can_be_constant_folded, copy_src_op_attributes, copy_tensor_attributes, extract_only_one_op, remove_tensor_from_sorted_graph, replace_tensor, sanitize_sorted_graph, toposort, IntVar, bmm_ccr, bmm_crr, bmm_rcr, bmm_rrr, permute021, unsqueeze, TensorAccessor, DimInfo, IntImm, Operator, Source, Tensor, gemm :autosummary: transform_special_ops ------------------------------------------- .. automodule:: aitemplate.compiler.transform.transform_special_ops :members: :imported-members: :exclude-members: gemm_rrr, is_singleton_dimension, gemm_rcr, gemm_rrr_small_nk, can_be_constant_folded, copy_src_op_attributes, copy_tensor_attributes, extract_only_one_op, remove_tensor_from_sorted_graph, replace_tensor, sanitize_sorted_graph, toposort, IntVar, bmm_ccr, bmm_crr, bmm_rcr, bmm_rrr, permute021, unsqueeze, TensorAccessor, DimInfo, IntImm, Operator, Source, Tensor, gemm :autosummary: transform_strided_op_and_view_op ------------------------------------------- .. automodule:: aitemplate.compiler.transform.transform_strided_op_and_view_op :members: :imported-members: :exclude-members: IntVar, is_singleton_dimension, DimInfo, IntImm, Operator, Source, Tensor, gemm :autosummary: transform_strided_ops ------------------------------------------- .. automodule:: aitemplate.compiler.transform.transform_strided_ops :members: :imported-members: :exclude-members: get_tensor_index, slice_reshape_scatter, slice_scatter, gen_tensor_index, IntVar, is_singleton_dimension, DimInfo, IntImm, Operator, Source, Tensor, gemm :autosummary: transform_strided_slice ------------------------------------------- .. automodule:: aitemplate.compiler.transform.transform_strided_slice :members: :imported-members: :exclude-members: dynamic_slice, get_tensor_index, slice_reshape_scatter, slice_scatter, gen_tensor_index, IntVar, is_singleton_dimension, DimInfo, IntImm, Operator, Source, Tensor, gemm :autosummary: --- ### Source Reference Utils (`docs/source/reference/utils.rst`) aitemplate.utils ================== visualization.plot ------------------ .. automodule:: aitemplate.utils.visualization.plot :members: :imported-members: :exclude-members: Tensor, Operator :autosummary: --- ### Source Runtime Cxx Design (`docs/source/runtime/cxx_design.rst`) ================ C++ Runtime Note ================ `Model` v.s. `ModelContainer` ============================= These are the two main classes involved in the C++ runtime implementation: * The bulk of the runtime implementation is in the `Model` class. * The `ModelContainer` class stores a set of shared constants and a collection of `Model` instances. Almost all functions in `model_interface.h` forward to a method in `ModelContainer`. When `Run` is invoked, `ModelContainer` looks for an available `Model`, or blocks until one becomes available (see the section on asynchronous predictions). It then forwards the run request to the runtime. Code Structure ============== Some important files: 1. `include/model_interface.h`: The interface that we expose in the compiled `.so`. 2. `include/model_container.h`: The bulk of the `ModelContainer` implementation. Some files are generated at compile time. These include: * `model-generated.h`: The implementation of the `Model`. * `model_container_base.cu`: A small part of the implementation for `ModelContainer` that needs to be generated. `ModelContainer` inherits from `ModelContainerBase`, and `ModelContainerBase`'s implementation lives in this file. See `model_container.h` for more details. All codegen templates can be found in `backend/main_templates.py`. The codegen implementation is in `backend/codegen.py`. Note that many of the headers in this directory rely on generated code and thus cannot be `#include` -d in external projects. `model_interface.h` is an exception. --- ### Source Runtime Index (`docs/source/runtime/index.rst`) Runtime Note ================== .. toctree:: :maxdepth: 1 cxx_design py_design --- ### Source Runtime Py Design (`docs/source/runtime/py_design.rst`) =================== Python Runtime Note =================== Python `Model` ============== `Model` is a collection of Python bindings to the C++ AIT runtime. This section describes the API. `AITData` --------- This class represents a contiguous blob of memory that AIT will use as a tensor. It is simply a named tuple with these fields: * `data_ptr: int`: An **unowned** pointer to **GPU** memory. In general, all of the APIs expect that this pointer will be valid for the entire duration of the call. * `shape: List[int]`: The shape of the tensor. * `dtype: str`: The tensor's dtype; one of `"float32", "float16", "int32", "int64"`. Note that most ops only support float16 at this stage. When using AITemplate with PyTorch, `AITData` can be constructed with the `torch_to_ait_data` utility: .. code-block:: python x = torch.randn(3, 3, 3).half().cuda() # Equivalent to AITData(x.data_ptr(), [3, 3, 3], "float16") x_ait = torch_to_ait_data(x) If PyTorch is not available, `Model` provides a set of functions for copying, allocating, and freeing GPU memory. See the docstrings in `compiler/model.py` for more information. `run` ----- `run` takes inputs and outputs as collections of `AITData` instances. Both arguments can be passed as either an ordered list or a dictionary (mapping name to tensor). .. code-block:: python # Arguments as a dictionary module.run( {"input0": in0_ait, "input1": in1_ait}, {"output0": out0_ait, "output1": out1_ait}, ) # Arguments as an ordered list. Note that you might need to query # the index mapping. input_name_to_idx = module.get_input_name_to_index_map() output_name_to_idx = module.get_output_name_to_index_map() inputs = [None] * len(input_name_to_idx) outputs = [None] * len(output_name_to_idx) for name in input_name_to_idx: inputs[input_name_to_idx[name]] = ait_inputs[name] for name in output_name_to_idx: outputs[output_name_to_idx[name]] = ait_outputs[name] module.run(inputs, outputs) One important caveat is that the output must have the **maximum** possible size. This is because of dynamic shapes: the size of the output may vary, but its shape is not inferred until inference time. The maximum shape can be queried with the `get_output_maximum_shape()`: .. code-block:: python # Can use either name or index. name_to_idx = module.get_output_name_to_idx() max_shape = module.get_output_maximum_shape(name_to_idx["output"]) max_shape = module.get_output_maximum_shape("output") `Model.run` returns a dictionary of output `AITData` instances with (possibly dynamic) shapes that inferred in the runtime. Nullptr Inputs/Outputs ---------------------- In general, inputs are allowed to be null if they are size 0 (e.g. at least one dimension is 0). The runtime enforces this with a check before any kernels are launched. .. code-block:: cpp If (input_name == nullptr && dim0 * dim1 * … * dimN != 0) { throw std::runtime_error(“input_name cannot be null!”); } This is convenient since torch.data_ptr() returns null for size zero tensors. The dynamic shape computation is skipped if the lower bound of the tensor’s size is positive. Constants --------- There are two types of constants in AIT; *bound* and *unbound* constants. A bound constant is known at compile time and may participate in constant folding. Bound constants are copied into GPU memory at model loading time. Values for bound constants may be provided by passing a dictionary (mapping constant name to AIT tensor) to `compile_model`. Unbound constants, on the other hand, do not participate in constant folding and must be provided before running the model. These must be set via `Model.set_constant`: .. code-block:: python module.set_constant("my_constant", AITData(...)) # The pointer in the the tensor must live for the entire duration of run() module.run(...) Constants are read-only and *shared* with all runtimes in the `ModelContainer`. `run_with_tensors` ------------------ `run_with_tensors` is a convenience method with the same interface as `run`, except it can take lists (or dicts) of `torch.Tensor` instances: .. code-block:: python input0 = torch.randn(input0_shape).cuda().half() output0 = torch.empty(output0_shape).cuda().half() # Returns a dictionary of reshaped outputs result = module.run_with_tensors([input0], [output0]) Streams and Asynchronous Predictions ------------------------------------ A pointer to a stream can optionally be passed to `run`. If none is given, the prediction happens on the default stream 0. If the `sync` argument is set to `True`, the stream is synchronized before `run()` returns. `sync` is `True` by default. Multiple predictions can happen at the same time (on the same or different streams). Under the hood, there is a fixed-size pool of runtime objects. When all the runtimes are used, `run()` blocks until one becomes available. The size of this pool can be configured with the `num_runtimes` option in `Model`'s constructor. CUDA Graph ---------- Run also takes a `graph_mode` option. If set to true, the runtime will try to use [CUDA graphs](https://developer.nvidia.com/blog/cuda-graphs/) to run the model. `graph_mode` is not supported on ROCm. The following is a high level overview of how graph mode works: 1) Each `Model` has an internal stream used for graph capturing. The model first runs all ops on this stream in capture mode. No kernel launches happen during this stage. 2) If this is the first run, a graph is instantiated via `cudaGraphInstantiate`. 3) On subsequent runs, we try to avoid the relatively expensive `cudaGraphInstantiate` call by updating the graph executor (`cudaGraphExecUpdate`). However, a new graph may still be instantiated if the topology of the graph somehow changed between runs. 4) Once we have the graph executor, we launch a single kernel on the stream that the user provided to `run()`. Graph mode is mainly beneficial when there are many small kernel launches. A lot of overhead can be avoided since there is only a single kernel launch in graph mode. --- ### Source Tutorial How To Add Op (`docs/source/tutorial/how_to_add_op.rst`) How to add an operator to the AIT codegen ========================================= This tutorial will demonstrate how to add a new operator to the AIT codegen. Full source code can be found at `examples/07_how_to_run_pt_model/how_to_run_pt_model.py`. 0. Prerequisites ---------------- We need to import necessary Python modules: .. code-block:: python from typing import Any, Dict, List import jinja2 import torch from aitemplate import backend from aitemplate.backend import registry from aitemplate.backend.backend_spec import CUDASpec, ROCMSpec from aitemplate.compiler import compile_model from aitemplate.compiler.base import IntVar, Operator, Tensor from aitemplate.testing import detect_target 1. Define the operator graph node --------------------------------- Graph nodes are usually defined at `aitemplate/compiler/ops`. .. code-block:: python class add_one(Operator): def __init__(self): super().__init__() # required, unique identity of operator category self._attrs["op"] = "add_one" # we can put whatever we want into the op attrs for later use self._attrs["has_profiler"] = False self._attrs["nop"] = False def __call__(self, x: Tensor) -> Tensor: # each operator needs to keep a record of input tensors self._attrs["inputs"] = [x] # optional, to set depth of the op based on inputs' depth, used in DFS self._set_depth() # infer output shape output_shape = self._infer_shape(x) # create output Tensor, of which the source op is the current op output = Tensor(output_shape, src_ops={self}) # remember current op's outputs self._attrs["outputs"] = [output] return output def _infer_shape(self, x) -> List[IntVar]: # infer output shape # In case of we need infer shape in C++ side, we will create a jinja2 template # for shape inference function, and render to Python code in graph node # and render the template into C++ code in codegen return x.shape() def gen_function(self) -> str: # this function will be used in codegen # here we only need to redirect to backend codegen function target = backend.target.Target.current() func_key = f"{target.name()}.{self._attrs['op']}.gen_function" func = registry.get(func_key) return func(self._attrs) .. note:: - `_attrs` in Operator is the most important data structure for codegen. - `_attrs["op"]` is the identity of operator category, which is used to find the corresponding codegen function in the backend; must be **unique**. 2. Define the necessary templates for Codegen --------------------------------------------- In AIT, there are 4 important templates for codegen: - `FUNC_TEMPLATE`: the template for generating the function body of the operator, and invoke GPU kernel in the body. - `FUNC_SIGNATURE_TEMPLATE`: the template for generating the function signature of the operator. The signature defines the name and arguments of the function. - `FUNC_CALL_TEMPLATE`: the template for generating the function call of the operator. The call will be used during inference to invoke the GPU kernel with given arguments. - `FUNC_DECL`: the template for forward declaration of the operator function. This is usually an alias of `FUNC_SIGNATURE_TEMPLATE`. .. code-block:: python FUNC_TEMPLATE = jinja2.Template( """ {{header_files}} namespace { {{kernel}} } // namespace {{func_signature}} { invoke_add_one(output, input, num_elements, stream); } """ ) FUNC_SIGNATURE = jinja2.Template( """ void {{func_name}}(half* output, const half* input, const int64_t num_elements, {{prefix}}Stream_t stream) """ ) FUNC_DECL = jinja2.Template( """ {{func_signature}}; """ ) FUNC_CALL_TEMPLATE = jinja2.Template( """ {{indent}}int64_t num_elements = 1; {% for dim_name in dim_names %} {{indent}}num_elements *= {{dim_name}}; {% endfor %} {{indent}}{{func_name}}( {{indent}} {{output}}, {{input}}, num_elements, stream /* default stream */ {{indent}}); """ ) 3. Create the GPU kernels ------------------------- In this example we use a simplest add one kernel. The kernel can be written by hand (as what programmer is expected to do), or generated by other tools. .. code-block:: python KERNEL_TEMPLATE = jinja2.Template( """ __global__ void add_one(half* output, const half* input, const int64_t num_elements) { const int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < num_elements) { output[idx] = input[idx] + half(1.0); } } void invoke_add_one(half* output, const half* input, int64_t num_elements, {{prefix}}Stream_t stream) { if (num_elements < 1024) { dim3 grid(1); dim3 block(num_elements); add_one<<>>(output, input, num_elements); } else { dim3 grid((num_elements + 1024 - 1) / 1024); dim3 block(1024); add_one<<>>(output, input, num_elements); } } """ ) (Optional) We also provide a helper function to handle CUDA/ROCm float16 data type difference. .. code-block:: python FUNC_CALL_FP16_PARAM_TEMPLATE = jinja2.Template( """reinterpret_cast( {% if is_cuda %}&({% endif %}{{name}}{% if is_cuda %}->raw()){% endif %})""" ) 4. Define the codegen function ------------------------------ The codegen function is the function that renders the templates we defined into valid C++ code string. The codegen function will take `func_attrs` from the graph node, and fill in the jinja2 template. .. code-block:: python def gen_function_call(func_attrs: Dict[str, Any], indent=" ", is_cuda=False) -> str: assert len(func_attrs["outputs"]) == 1 assert len(func_attrs["inputs"]) == 1 output_name = FUNC_CALL_FP16_PARAM_TEMPLATE.render( name=func_attrs["outputs"][0]._attrs["name"], is_cuda=is_cuda ) input_name = FUNC_CALL_FP16_PARAM_TEMPLATE.render( name=func_attrs["inputs"][0]._attrs["name"], is_cuda=is_cuda ) dim_names = [dim._attrs["name"] for dim in func_attrs["inputs"][0].shape()] return FUNC_CALL_TEMPLATE.render( func_name=func_attrs["name"], output=output_name, input=input_name, dim_names=dim_names, indent=indent, ) def gen_function(func_attrs: Dict[str, Any], header_files: str, backend_spec) -> str: prefix = backend_spec.prefix return FUNC_TEMPLATE.render( header_files=header_files, kernel=KERNEL_TEMPLATE.render(prefix=prefix), func_signature=FUNC_SIGNATURE.render( func_name=func_attrs["name"], prefix=prefix ), ) def gen_function_decl(func_attrs: Dict[str, Any], backend_spec) -> str: return FUNC_DECL.render( func_signature=FUNC_SIGNATURE.render( func_name=func_attrs["name"], prefix=backend_spec.prefix, ).strip() ) 5.1 Register the codegen function in CUDA backend ------------------------------------------------- CUDA backend functions are usually defined at `aitemplate/backend/cuda/`. .. code-block:: python CUDA_HEADER_FILES = """ #include """ @registry.reg("cuda.add_one.gen_function") def cuda_add_one_gen_function(func_attrs: Dict[str, Any]) -> str: return gen_function(func_attrs, CUDA_HEADER_FILES, CUDASpec()) @registry.reg("cuda.add_one.func_decl") def cuda_add_one_gen_function_decl(func_attrs: Dict[str, Any]) -> str: return gen_function_decl(func_attrs, CUDASpec()) @registry.reg("cuda.add_one.func_call") def cuda_add_one_gen_function_call(func_attrs: Dict[str, Any], indent=" ") -> str: return gen_function_call(func_attrs, indent, is_cuda=True) 5.2 (Optional) Register the codegen function to ROCm backend ------------------------------------------------------------ ROCm backend functions are usually defined at `aitemplate/backend/rocm/`. .. code-block:: python HIP_HEADER_FILES = """ #include #include """ @registry.reg("rocm.add_one.gen_function") def rocm_add_one_gen_function(func_attrs: Dict[str, Any]) -> str: return gen_function(func_attrs, HIP_HEADER_FILES, ROCMSpec()) @registry.reg("rocm.add_one.func_decl") def rocm_add_one_gen_function_decl(func_attrs: Dict[str, Any]) -> str: return gen_function_decl(func_attrs, ROCMSpec()) @registry.reg("rocm.add_one.func_call") def rocm_add_one_gen_function_call(func_attrs: Dict[str, Any], indent=" ") -> str: return gen_function_call(func_attrs, indent, is_cuda=False) 6. Compile and verify the results with PyTorch ---------------------------------------------- .. code-block:: python def create_ait_model(shapes): X = Tensor( shape=shapes, dtype="float16", name="X", is_input=True, ) Y = add_one()(X) Y._attrs["is_output"] = True Y._attrs["name"] = "Y" return Y def verify_add_one(): shapes = [16, 512] x = torch.randn(shapes).cuda().half() y_pt = x + 1.0 Y = create_ait_model([16, 512]) target = detect_target() with compile_model(Y, target, "./tmp", "add_one") as module: y = torch.empty(shapes).cuda().half() inputs = {"X": x} outputs = {"Y": y} module.run_with_tensors(inputs, outputs) print(torch.allclose(y, y_pt, atol=1e-2, rtol=1e-2)) --- ### Source Tutorial How To Infer Pt (`docs/source/tutorial/how_to_infer_pt.rst`) How to inference a PyTorch model with AIT ========================================= This tutorial will demonstrate how to inference a PyTorch model with AIT. Full source code can be found at `examples/07_how_to_run_pt_model/how_to_run_pt_model.py`. 0. Prerequisites ---------------- We need to import necessary Python modules: .. code-block:: python import torch from aitemplate.compiler import compile_model from aitemplate.frontend import nn, Tensor from aitemplate.testing import detect_target from aitemplate.testing.benchmark_pt import benchmark_torch_function from aitemplate.utils.graph_utils import sorted_graph_pseudo_code 1. Define a PyTorch module -------------------------- Here we define a PyTorch model which is commonly seen in Transformers: .. code-block:: python class PTSimpleModel(torch.nn.Module): def __init__(self, hidden, eps: float = 1e-5): super().__init__() self.dense1 = torch.nn.Linear(hidden, 4 * hidden) self.act1 = torch.nn.functional.gelu self.dense2 = torch.nn.Linear(4 * hidden, hidden) self.layernorm = torch.nn.LayerNorm(hidden, eps=eps) def forward(self, input): hidden_states = self.dense1(input) hidden_states = self.act1(hidden_states) hidden_states = self.dense2(hidden_states) hidden_states = hidden_states + input hidden_states = self.layernorm(hidden_states) return hidden_states 2. Define an AIT module ----------------------- We can define a similar AIT module as follows: .. code-block:: python class AITSimpleModel(nn.Module): def __init__(self, hidden, eps: float = 1e-5): super().__init__() self.dense1 = nn.Linear(hidden, 4 * hidden, specialization="fast_gelu") self.dense2 = nn.Linear(4 * hidden, hidden) self.layernorm = nn.LayerNorm(hidden, eps=eps) def forward(self, input): hidden_states = self.dense1(input) hidden_states = self.dense2(hidden_states) hidden_states = hidden_states + input hidden_states = self.layernorm(hidden_states) return hidden_states .. warning:: The `nn.Module` API in AIT looks similar to PyTorch, but it is not the same. The fundamental difference is that AIT module is a container to build a graph, while PyTorch module is a container to store parameters for eager. Which means, each AIT module's `forward` method can be only called once, and the graph is built during the first call. If you want to share parameters, you need to use the `compiler.ops` instead. The `compiler.ops` is similar to `functional` in PyTorch. AITemplate supports automatic fusion of linear followed by other operators. However in many cases, especially for quick iterations, we use manual `specialization` to specify the fused operator. For example, `specialization="fast_gelu"` will fuse linear with the `fast_gelu` operator. 3. Define a helper function to map PyTorch parameters to AIT parameters ----------------------------------------------------------------------- In AIT, all names must follow the C variable naming standard, because the names will be used in the codegen process. .. code-block:: python def map_pt_params(ait_model, pt_model): ait_model.name_parameter_tensor() pt_params = dict(pt_model.named_parameters()) mapped_pt_params = {} for name, _ in ait_model.named_parameters(): ait_name = name.replace(".", "_") assert name in pt_params mapped_pt_params[ait_name] = pt_params[name] return mapped_pt_params .. warning:: - Different to PyTorch, it is required to call ait_model **.name_parameter_tensor()** method to provide each parameter with a name with a direct map to PyTorch. - Because all names in AIT must follow the C variable naming standard, you can easily replace `.` by `_` or use a regular expression to make sure the name in valid. - For networks with conv + bn subgraph, we currently don't provide an automatic pass to fold it. Please refer to our ResNet and Detectron2 examples to see how we handle CNN layout transform and BatchNorm folding. 4. Create PyTorch module, inputs/outputs ---------------------------------------- .. code-block:: python batch_size=1024 hidden=512 # create pt model pt_model = PTSimpleModel(hidden).cuda().half() # create pt input x = torch.randn([batch_size, hidden]).cuda().half() # run pt model pt_model.eval() y_pt = pt_model(x) 5. Create AIT module, inputs/outputs ------------------------------------ .. code-block:: python batch_size=1024 hidden=512 # create AIT model ait_model = AITSimpleModel(hidden) # create AIT input Tensor X = Tensor( shape=[batch_size, hidden], name="X", dtype="float16", is_input=True, ) # run AIT module to generate output tensor Y = ait_model(X) # mark the output tensor Y._attrs["is_output"] = True Y._attrs["name"] = "Y" .. warning:: - Similar to MetaTensor, LazyTensor and a lot of other lazy evaluation frameworks, AIT's Tensor records the computation graph, and the graph is built when the Tensor is compiled. - For input tensor, it is required to set the attribute **is_input=True**. - For output tensor, it is required to set the attribute **Y._attrs["is_output"] = True**. - For input and output tensors, it is better to provide the **name** attributes to use in runtime. 6. Compile AIT module into runtime and do verification ------------------------------------------------------ .. code-block:: python # map pt weights to ait weights = map_pt_params(ait_model, pt_model) # codegen target = detect_target() with compile_model( Y, target, "./tmp", "simple_model_demo", constants=weights ) as module: # create storage for output tensor y = torch.empty([batch_size, hidden]).cuda().half() # inputs and outputs dict inputs = {"X": x} outputs = {"Y": y} # run module.run_with_tensors(inputs, outputs, graph_mode=True) # verify output is correct print(torch.allclose(y, y_pt, atol=1e-2, rtol=1e-2)) # benchmark ait and pt count = 1000 ait_t, _, _ = module.benchmark_with_tensors( inputs, outputs, graph_mode=True, count=count ) print(f"AITemplate time: {ait_t} ms/iter") pt_t = benchmark_torch_function(count, pt_model.forward, x) print(f"PyTorch eager time: {pt_t} ms/iter") In this example, AIT will automatically fuse GELU and elementwise addition into the TensorCore/MatrixCore gemm operation. On RTX-3080, in the example AIT is about 1.15X faster than PyTorch Eager. .. note:: - In this example, we fold the parameters (`weights`) into AIT runtime. The final dynamic library will contain them as parameters. - If during the compile time we don't provide the parameters (for example, because the total parameters size is greater than 2GB), we can always call `set_constant` function in the runtime. Please check the runtime API for the details. --- ### Source Tutorial How To Visualize (`docs/source/tutorial/how_to_visualize.rst`) How to visualize an AIT model ============================= Visualization is important for understanding the behavior of a model optimization. In AIT, we modify the codegen a little bit, from generating CUDA/HIP C++ code to HTML/Javascript code, then we can generate a visualization of the model. The following code will generate a visualization of our first example. 1. Define the AIT Model ----------------------- .. code-block:: python from aitemplate import compiler from aitemplate.frontend import nn, Tensor from aitemplate.testing import detect_target from aitemplate.utils.visualization import plot_graph class AITSimpleModel(nn.Module): def __init__(self, hidden, eps: float = 1e-5): super().__init__() self.dense1 = nn.Linear(hidden, 4 * hidden, specialization="fast_gelu") self.dense2 = nn.Linear(4 * hidden, hidden) self.layernorm = nn.LayerNorm(hidden, eps=eps) def forward(self, input): hidden_states = self.dense1(input) hidden_states = self.dense2(hidden_states) hidden_states = hidden_states + input hidden_states = self.layernorm(hidden_states) return hidden_states def gen_ait_model(): batch_size = 512 hidden = 1024 ait_model = AITSimpleModel(hidden) ait_model.name_parameter_tensor() X = Tensor( shape=[batch_size, hidden], name="X", dtype="float16", is_input=True, ) Y = ait_model(X) Y._attrs["is_output"] = True Y._attrs["name"] = "Y" return Y output_tensor = gen_ait_model() 2. Apply optimizations on the AIT Model --------------------------------------- .. code-block:: python def apply_optimizations(tensors): target = detect_target() # first, convert output tensors to graph with target: graph = compiler.transform.toposort(tensors) # second, provide names to the graph compiler.transform.name_graph(graph) compiler.transform.mark_param_tensor(graph) compiler.transform.mark_special_views(graph) # we can apply optimizations to the graph, or test single optimization pass on the graph graph = compiler.transform.optimize_graph(graph, "./tmp") return graph graph = apply_optimizations(output_tensor) 3. Generate visualization ------------------------- .. code-block:: python # Plot the graph plot_graph(graph, file_path="ait_model.html") The visualization will be generated in the "ait_model.html" file. This file can be opened in Chrome without any web server. .. raw:: html --- ### Source Tutorial Index (`docs/source/tutorial/index.rst`) Tutorials ========= .. toctree:: :maxdepth: 1 how_to_infer_pt how_to_add_op how_to_visualize --- --- METRICS --- - Files Extracted: 38 - Estimated Token Budget: ~33743 tokens - Canonical Reference: https://codewiki.google/github.com/facebookincubator/AITemplate