This contains the main codebases of the paper "Lightening-Transformer: A Dynamically-operated Optically-interconnected Photonic Transformer Accelerator".
Accepted by HPCA 2024.
We provides three kinds of codebases:
-
(1) algorithm codes for training/running models on our photonic accelerator, with the analytic transformation of our unique photonic tensor core embedded in the computation process. See
./software_modelfor detailed implementation and usages, including the DeiT case. -
(2) hardware simulator for estimating the energy and latency running Transformers on our photonic accelerator. See
./hardware_simulatorfor detailed implementation and usages. -
(3) profile codes for profiling latency and power usage of running Transformers on GPU. See
./profilefor detailed implementation and usages. The implementation refers to Neurips'23, Pre-RMSNorm and Pre-CRMSNorm Transformers: Equivalent and Efficient Pre-LN Transformers.
The DeiT requires to install PyTorch and torchvision 0.8.1+ and pytorch-image-models 0.3.2.
conda create -n test # create a virtual env
conda install pytorch torchvision torchaudio pytorch-cuda=your_cuda_version -c pytorch -c nvidia # install pytorch
pip install timm==0.3.2 torchpack packaging einops gdown
conda activate test # activate the test env
For torch.2.0+, you will encounter the ModuleNotFoundError: No module named 'torch._six' in '/path_to_your_conda_envs/your_env_name/lib/python_version/site-packages/timm/models/layers/helpers.py". This is because torch2.0 doesn't have torch._six. Please replace the helper.py file with the following one.
from itertools import repeat
# from torch._six import container_abcs
# From PyTorch internals
def _ntuple(n):
def parse(x):
if isinstance(x, str):
return x
return tuple(repeat(x, n))
return parse
to_1tuple = _ntuple(1)
to_2tuple = _ntuple(2)
to_3tuple = _ntuple(3)
to_4tuple = _ntuple(4)
to_ntuple = _ntuple
[1] Hanqing Zhu, Jiaqi Gu, Hanrui Wang, Zixuan Jiang, Rongxing Tang, Zhekai Zhang, Chenghao Feng, Song Han, Ray T. Chen and David Z. Pan. "Lightening-Transformer: A Dynamically-operated Optically-interconnected Photonic Transformer Accelerator", IEEE International Symposium on High-Performance Computer Architecture (HPCA'24).