DiJiang - GitLite

KO | EN

gitlite — search

#hacktoberfest #typescript #llm #go #android #linux #macos #java #awesome #vue #pytorch #ios #mcp #openai

DiJiang

(★ 103)

[ICML'24 Oral] The official code of "DiJiang: Efficient Large Language Models through Compact Kernelization", a novel DCT-based linear attention mechanism.

File Explorer

Download Latest Version (.zip)

- dijiang.png
- dijiang_speed.png
- experiments1.png
- experiments2.png
- scheme.png
- - config.json
  - modeling_gpt_neox_dijiang.py
  - special_tokens_map.json
  - tokenizer.json
  - tokenizer_config.json
- - config.json
  - modeling_gpt_neox_dijiang.py
  - special_tokens_map.json
  - tokenizer.json
  - tokenizer_config.json
- - config.json
  - modeling_gpt_neox_dijiang.py
  - special_tokens_map.json
  - tokenizer.json
  - tokenizer_config.json
- - config.json
  - modeling_gpt_neox_dijiang.py
  - special_tokens_map.json
  - tokenizer.json
  - tokenizer_config.json
- - config.json
  - modeling_gpt_neox_dijiang.py
  - special_tokens_map.json
  - tokenizer.json
  - tokenizer_config.json
README.md
requirements.txt
run_dijiang.sh
train_pythia.py

# Installation Guide

1. Get the code

git clone https://github.com/YuchuanTian/DiJiang

Downloads the entire project code from GitHub to your computer.

cd DiJiang

Moves into the project folder you just downloaded.

2. Python

Easy Recommended

Prerequisites

Git Needed to download the project code from GitHub.
Python 3 On Windows, be sure to check 'Add Python to PATH' during installation.

pip install -r requirements.txt

Installs the Python libraries listed in requirements.txt (or similar).

python <실행할 파일명>.py # README에서 정확한 실행 파일명을 확인하세요

Runs the Python script (or module).

✅ If it runs without errors and prints output in the terminal, it worked.

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# DiJiang: Efficient Large Language Models through Compact Kernelization

<p align="left">
<a href="https://arxiv.org/abs/2403.19928" alt="arXiv">
    <img src="https://img.shields.io/badge/arXiv-2403.19928-b31b1b.svg?style=flat" /></a>
</p>

## News

**6/14/2024** Our work DiJiang is accepted as **ICML'24 Oral**!🎉🎉🎉

<img src="./imgs/dijiang.png" alt="scheme" style="zoom: 20%;"/>

This is the official implementation of ["DiJiang: Efficient Large Language Models through Compact Kernelization"](https://arxiv.org/abs/2403.19928), a novel Frequency Domain Kernelization approach to linear attention. DiJiang is a mythical monster in Chinese tradition. It has six legs and four wings , and it is known for its capability to move quickly.

![scheme](https://raw.githubusercontent.com/YuchuanTian/DiJiang/main/imgs/scheme.png)

🚀 In our work "DiJiang", we propose a fast linear attention approach that leverages the power of **Fast Discrete Cosine Transform (Fast DCT).** Theoretical analyses and extensive experiments demonstrate the speed advantage of our method.

<img src="./imgs/dijiang_speed.png" alt="scheme" style="zoom: 35%;" />

<img src="./imgs/experiments1.png" alt="scheme" style="zoom: 67%;" />

<img src="./imgs/experiments2.png" alt="scheme" style="zoom: 67%;" />

## Training

### Preparation

- Download Dataset: We utilized the Pile dataset for training [The Pile: An 800GB Dataset of Diverse Text for Language Modeling (arxiv.org)](https://arxiv.org/abs/2101.00027)
- Check CUDA version: suggested cuda>=11.7

- Install required packages: ```pip install -r requirements.txt```

- Download pre=trained model weights (i.e. ```pytorch_model.bin``` files in the original HuggingFace models) and place them in the corresponding DiJiang model folder (```e.g. modeling/pythia-70M-dijiang```).

Links of pre-trained weights:

```
# Pythia-70M
https://huggingface.co/EleutherAI/pythia-70m/resolve/main/pytorch_model.bin?download=true

# Pythia-160M
https://huggingface.co/EleutherAI/pythia-160m/resolve/main/pytorch_model.bin?download=true

# Pythia-410M
https://huggingface.co/EleutherAI/pythia-410m/resolve/main/pytorch_model.bin?download=true

# Pythia-1B
https://huggingface.co/EleutherAI/pythia-1b/resolve/main/pytorch_model.bin?download=true

# Pythia-2.8B
https://huggingface.co/EleutherAI/pythia-2.8b/resolve/main/pytorch_model.bin?download=true

```

### Running

A training shell is provided at ```bash run_dijiang.sh``` for your reference.

Important arguments to take care of:

```

--model_name_or_path <path to the HuggingFace model>
    --dataset_dir <path to the original PILE dataset>
    --data_cache_dir <path to the cached dataset (if cached); cache stored at this dir>
    --per_device_train_batch_size <global_batchsize // device_count>
    --output_dir <dir to store the trained model>
    --logging_dir <dir to store tensorboard>

```

## Future work (Stay Tuned!)
- [x] Training code for **DiJiang**
- [ ] Better inference code
- [ ] Inference models

## Citation (in BibTex)

```
@misc{chen2024dijiang,
      title={DiJiang: Efficient Large Language Models through Compact Kernelization}, 
      author={Hanting Chen and Zhicheng Liu and Xutao Wang and Yuchuan Tian and Yunhe Wang},
      year={2024},
      eprint={2403.19928},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}
```

## Acknowledgement

We acknowledge the authors of the following repos:

[DenseSSM (github)](https://github.com/WailordHe/DenseSSM)

[Llama2 (huggingface.co)](https://huggingface.co/docs/transformers/main/model_doc/llama2#usage-tips)

[Pythia (EleutherAI)](https://huggingface.co/docs/transformers/main/model_doc/llama2#usage-tips)

[RetNet (Microsoft)](https://github.com/microsoft/torchscale/tree/main)

[Transformersscale (Microsoft)](https://github.com/microsoft/transformersscale/tree/main)

[Stanford Alpaca (Stanford)](https://github.com/tatsu-lab/stanford_alpaca)

[Chinese-LLaMA-Alpaca (ymcui)](https://github.com/ymcui/Chinese-LLaMA-Alpaca)