Sparse Transformer

This repository is an optimized version of the official SparseTransformer codebase, featuring FP16 and BF16 support to significantly reduce GPU memory usage.

1. Environment

• CUDA：11.8
• PyTorch：2.0.0

2. Install

• install with pip

export TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX" # support different GPU architectures [optional]
export ENABLE_BF16=1 # support BF16 [optional]
python3 -m pip install SparseTransformer

• install from source

git clone https://github.com/lh9171338/SparseTransformer.git
cd SparseTransformer

export TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX" # support different GPU architectures [optional]
export ENABLE_BF16=1 # support BF16 [optional]
python3 setup.py install

3. Usage

Please refer to the official README.md for usage instructions.

4. Performance Evaluation

4.1 Operator Benchmarking

Tests were conducted on two attention computation modes:

• Without relative position encoding (using attention_step1 and attention_step2 operators)
• With relative position encoding (using dot_prod_with_idx_all and attention_step2_with_rel_pos_value operators)

Input voxel count: 268,720

(1) Memory Consumption Comparison

FP16 and BF16 showed identical memory usage, reducing consumption by 33% and 35% respectively compared to FP32.

(2) Precision Comparison

Maximum errors in output and gradients were measured

• Without Relative Position Encoding

• With Relative Position Encoding

4.2 Semantic Segmentation Task Evaluation

Tests conducted using SphereFormer model

(1) Training Mode

With configuration parameters held constant except for the SpTr module's dtype, FP16/BF16 implementations demonstrate equivalent mIoU accuracy while reducing memory consumption by 31% (20.3GB => 14.1GB)

(2) Inference Mode

Same model evaluated with different dtype, both FP16 and BF16 significantly reduce memory usage by 31% (18.1GB => 15.4GB)