가산 양자화(Additive Quantization)를 통한 대규모 언어 모델(LLM)의 극한 압축을 위한 공식 PyTorch 구현체

가산 양자화 (Additive Quantization)를 통한 대규모 언어 모델 (Large Language Models)의 극한 압축을 위한 공식 PyTorch 구현체

[2025.04] aqlm v1.1.7을 출시했습니다. GPU에서 임의의 8차원 코드북 (codebooks) 지원을 추가하였으며, 1비트 모델의 정확도를 개선했습니다. 예를 들어, ISTA-DASLab/Llama-2-7b-AQLM-1Bit-1x8-hf 모델은 약 1비트에서 WikiText 2 PPL 7.85를 달성합니다. 자신의 모델을 이 방식으로 양자화하려면 아래 튜토리얼에 따라 num_codebooks=1, nbits_per_codebook=256을 사용하세요.

[2024.11] PV-tuning이 NeurIPS'2024 오럴 발표(oral presentation)로 채택되었습니다!

[2024.05] AQLM이 ICML'2024에 채택되었습니다! 컨퍼런스에 참석하신다면 이 포스터 근처에서 저희를 만나보세요.

[2024.06] PV-tuning이라는 새로운 미세 조정 (finetuning) 알고리즘으로 AQLM을 확장하는 새로운 논문을 발표했습니다. 또한 이 컬렉션에 PV-tuning된 AQLM 모델들을 공개합니다.

[2024.08] PV-Tuning 브랜치를 메인 브랜치에 병합했습니다. 이전 미세 조정(8월 21일 이전) 결과들을 재현하려면 커밋 559a366을 사용하세요.

다음 Google Colab 예제를 사용하여 사전 양자화된 (prequantized) 모델을 실행하는 방법을 알아보세요:

| 기본 AQLM 생성 |
| GPU/CPU 스트리밍 |
| CUDA 그래프를 이용한 추론 (3배 속도 향상) |
| PEFT를 이용한 미세 조정 (Fine-tuning) |
| vLLM을 이용한 서빙 |
|---|---|---|---|

브라우저에서 직접 CPU로 AQLM+PV 추론을 시도해보고 싶다면 aqlm-rs를 확인하세요:

라이브 데모 (Live demo): galqiwi.github.io/aqlm-rs/about.html
소스 코드 (Source code): galqiwi/demo-aqlm-rs

이 저장소는 현재 LLaMA, Mistral 및 Mixtral 제품군 모델과 작동하도록 설계되었습니다.
아래에 보고된 모델들은 부록 A에 설명된 대로, 교사 로짓 (teacher logits)을 사용한 교차 엔트로피 (cross-entropy) 목적 함수와 함께 **전체 모델 미세 조정 (full model fine-tuning)**을 사용했습니다.

PV-Tuning이 적용되지 않은 여러 사전 양자화된 AQLM 모델을 제공합니다 (PV-Tuned 모델은 아래로 스크롤하세요):

모델 (Model)	AQLM 방식 (scheme)	WikiText-2 PPL	MMLU (5-shot) FP16→AQLM	모델 크기 (Gb)	Hub 링크
Llama-3-8b	1x16	-	0.65→0.56	4.1	Link
...

또한 PV-tuning을 통해 조정된 AQLM 모델을 다운로드할 수 있습니다:

Model	AQLM scheme	WikiText-2 PPL	Model size, Gb	Hub link
Llama-2-7b	1x16g8	5.68	2.4	Link
...	7.85	1.34	Link
Llama-2-13b	1x16g8	5.05	4.1	Link
Llama-2-70b	1x16g8	3.78	18.8	Link
...
Note that models with "g16" in their scheme require aqlm inference library v1.1.6 or newer:

pip install aqlm[gpu,cpu]>=1.1.6

Above perplexity is evaluated on 4k context length for Llama 2 models and 8k for Mistral/Mixtral and Llama 3.
Please also note that token-level perplexity can only be compared within the same model family, but should not be compared between models that use different vocabularies.
While Mistral has a lower perplexity than Llama 3 8B but this does not mean that Mistral is better: Llama's perplexity is computed on a much larger dictionary and has higher per-token perplexity because of that.

For more evaluation results and detailed explanations, please see our papers: Egiazarian et al. (2024) for pure AQLM and Malinovskii et al. (2024) for PV-Tuned models.

AQLM quantization setpus vary mainly on the number of codebooks used as well as the codebook sizes in bits. The most popular setups, as well as inference kernels they support are:

Kernel	Number of codebooks	Codebook size, bits	Scheme Notation	Accuracy	Speedup	Fast GPU inference	Fast CPU inference
Triton	K	N	KxN	-	Up to ~0.7x	✅	❌
...
To run the models, one would have to install an inference library:

pip install aqlm[gpu,cpu]

, specifying either gpu

, cpu
or both based on one's inference setting.

Then, one can use the familiar .from_pretrained

method provided by the transformers library:

from transformers import AutoModelForCausalLM
quantized_model = AutoModelForCausalLM.from_pretrained(
"ISTA-DASLab/Llama-2-7b-AQLM-2Bit-1x16-hf",
...

Notice that torch_dtype
should be set to either torch.float16
or "auto"

GPU에서는 torch.float16 또는 "auto"로, CPU에서는 torch.float32로 torch_dtype을 설정해야 합니다. 그 이후에는 양자화되지 않은 모델과 완전히 동일하게 모델을 사용할 수 있습니다.

requirements.txt로부터 패키지를 설치하세요.

:

pip install -r requirements.txt

스크립트는 관련 토크나이저(tokenizer)와 데이터셋(datasets)을 다운로드하여 로컬에 캐싱해야 합니다. 환경 변수에 의해 대체 경로가 제공되지 않는 한, 이들은 기본 Huggingface Datasets 디렉토리에 저장됩니다. 관련 Datasets 문서 섹션을 참조하세요.

AQLM으로 모델을 양자화할 때는 모델이 학습되었던 원본 데이터의 하위 집합(subset)을 사용하는 것을 권장합니다.

Llama-2 모델의 경우, 사용 가능한 가장 유사한 데이터셋은 RedPajama입니다. RedPajama의 하위 집합을 로드하려면 --dataset 인자에 "pajama"를 제공하세요. 이는 nsamples 데이터를 처리하고 제공된 모델 토크나이저를 사용하여 토큰화(tokenize)합니다.

또한, 우리는 4096 컨텍스트 길이(context lengths)에 대해 Llama 및 Solar/Mistral 모델용으로 토큰화된 Redpajama를 Huggingface에 저장하여 제공합니다. 이를 로드하려면 다음을 사용하세요:

from huggingface_hub import hf_hub_download
hf_hub_download(repo_id="Vahe1994/AQLM", filename="data/name.pth", repo_type="dataset")

HF에서 다운로드한 데이터를 사용하려면, 이를 data 폴더에 넣고(선택 사항) main.py의 "--dataset" 인자에 올바른 경로를 설정하세요.

경고: 이 하위 집합들은 이미 해당 모델 토크나이저로 처리되었습니다. 만약 다른 모델(예: mistral/mixtral)을 양자화하고 싶다면, src/datautils에 제공된 스크립트를 사용하여 데이터를 다시 토큰화하십시오.

선택적으로 데이터를 Weights and Biases 서비스 (wandb)에 로깅할 수 있습니다.

W&B 로깅을 위해 pip install wandb를 실행하세요.

실험을 실행하기 전에 $WANDB_ENTITY, $WANDB_PROJECT, $WANDB_NAME 환경 변수를 지정하십시오. 로깅을 활성화하려면 --wandb 인자를 사용하세요.

이 코드는 여러 대의 80GB GPU RAM을 갖춘 A100 GPU를 사용하여 개발 및 테스트되었습니다.
VRAM 사용량을 줄이려면 --offload activations 옵션을 사용할 수 있습니다.
Language Model Evaluation Harness용

평가를 수행하려면 하나 또는 여러 개의 장치에 전체 모델과 활성화 텐서 (activation tensors)를 로드할 수 있는 충분한 메모리가 있어야 합니다.

AQLM 양자화 (quantization)는 GPTQ와 같은 더 단순한 양자화 방법보다 보정 (calibrate)에 훨씬 더 오랜 시간이 걸립니다. 이는 양자화 시간에만 영향을 미치며, 추론 시간 (inference time)에는 영향을 주지 않습니다.

예를 들어, 기본 설정으로 7B 모델을 양자화하는 데 단일 A100 GPU에서 약 1일이 소요됩니다. 마찬가지로, 단일 GPU에서 70B 모델을 양자화하는 데는 10~14일이 걸릴 것입니다. 빠른 상호 연결 (interconnect) 기능을 갖춘 여러 개의 GPU가 있다면, 단순히 여러 GPU에 대해 CUDA_VISIBLE_DEVICES를 설정함으로써 AQLM 멀티 GPU (multi-gpu)를 실행하여 속도를 높일 수 있습니다. 두 개의 GPU에서 7B 모델을 양자화하면 양자화 시간이 약 14.5시간으로 단축됩니다. 마찬가지로, 8개의 A100 GPU에서 70B 모델을 양자화하는 데는 3일 18시간이 소요됩니다.

GPU를 추가하지 않고 양자화 속도를 높여야 한다면, --relative_mse_tolerance를 높이거나

--init_max_points_per_centroid를 설정하거나

--finetune_max_epochs를 제한할 수도 있습니다.
하지만 이는 대개 모델 정확도 (accuracy) 저하라는 대가를 치르게 됩니다.

코드를 실행하려면 LLaMA 모델이 Huggingface 형식으로 다운로드되어 로컬에 저장되어 있어야 합니다. 아래 스크립트들은 $TRANSFORMERS_CACHE 변수가 Huggingface Transformers 캐시 폴더를 가리키고 있다고 가정합니다.
모델을 다운로드하고 캐싱하려면 동일한 환경에서 다음을 실행하십시오:

from transformers import AutoTokenizer, AutoModelForCausalLM
model_name = "meta-llama/Llama-2-7b-hf" # 또는 다운로드하고자 하는 다른 모델명
tokenizer = AutoTokenizer.from_pretrained(model_name, torch_dtype="auto")
...

이 스크립트는 모델을 압축한 다음, WikiText2, C4, Penn Treebank 데이터셋을 사용하여 퍼플렉시티 (perplexity) 측면에서 성능을 테스트합니다.

스크립트를 실행하는 명령은 다음과 같습니다:

export CUDA_VISIBLE_DEVICES=0 # 또는 예: 0,1,2,3
export MODEL_PATH=<HUB에 있는 모델 경로>
export DATASET_PATH=<데이터셋 이름 또는 커스텀 데이터 경로 입력>
...

주요 CLI 인자 (arguments):

CUDA_VISIBLE_DEVICES

• 기본적으로 코드는 사용 가능한 모든 GPU를 사용합니다. 특정 GPU(또는 단일 GPU)를 사용하려면 이 변수를 사용하세요.

MODEL_PATH

• Hugging Face 허브(예: meta-llama/Llama-2-7b-hf) 또는 transformers 모델과 토크나이저(tokenizer)가 포함된 로컬 폴더 경로입니다.

DATASET_PATH

• 보정(calibration) 데이터 경로(위 내용 참조) 또는 표준 데이터셋[c4, ptb, wikitext2]입니다.

• Llama-2 모델의 경우, 다음을 사용할 수 있습니다:
  DATASET_PATH=./data/red_pajama_n=1024_4096_context_length.pth

RedPajama의 일부(최대 1024개 샘플)를 위한 경로입니다.

• Llama-2 모델의 경우, 다음을 사용할 수 있습니다:

--nsamples

• 보정 데이터 시퀀스(sequences)(훈련 + 검증)의 수입니다. 이 파라미터가 설정되지 않으면 사용 가능한 모든 보정 데이터를 사용합니다.

--val_size

• 블록 미세 조정(block finetuning) 시 조기 종료(early stopping)를 위한 검증 시퀀스 수입니다. 기본값은 0이며, --nsamples보다 작아야 합니다.

--num_codebooks

• 레이어당 코드북(codebook)의 수입니다.

--nbits_per_codebook

• 각 코드북은 2 ** nbits_per_codebook 개의 벡터를 포함합니다.

--in_group_size

• 함께 양자화되는 가중치의 수 (arXiv 논문의 "g"와 동일)

--finetune_batch_size

• (미세 조정 시에만 해당) 각 최적화 단계(optimization step)에서 사용되는 총 시퀀스 수입니다.

--local_batch_size

• finetune_batch_size를 누적할 때, GPU당 순전파(forward pass) 한 번에 처리할 샘플 수입니다 (GPU RAM 사용량에 영향을 미침).

--relative_mse_tolerance

• (초기 보정 시) (현재 에포크 MSE / 이전 에포크 MSE) > (1 - relative_mse_tolerance)가 되면 훈련을 중단합니다.

--finetune_max_epochs

• 블록 튜닝(block tuning) 시 보정 데이터를 통과하는 최대 횟수입니다.

--finetune_early_stop

• 검증 데이터에서 개선이 없을 때 보정 데이터를 통과하는 최대 횟수입니다.

--offload_activations

• 보정 중에 활성화(activations) 값을 GPU 메모리에서 RAM으로 이동합니다. 이는 VRAM 사용량을 줄이지만 (하드웨어에 따라) 보정 속도를 약 10% 느리게 합니다.

--save

• 양자화된 모델을 저장/로드할 경로입니다. (--load 도 참조하세요)

--wandb

• 이 파라미터가 설정되면, 코드는 결과를 wandb에 기록합니다.

--attn_implementation

• 어텐션 (attention) 방식을 지정합니다 (transformers >= 4.38 이상).
  Sdpa 어텐션은 때때로 문제를 일으킬 수 있으므로 eager 구현체를 사용하는 것을 권장합니다.

추가적인 하이퍼파라미터(hyperparameters)를 사용할 수 있습니다. 압축 파라미터를 포함한 명령줄 인자(command line arguments)에 대한 자세한 내용은 python main.py --help를 실행하여 확인하세요.

이 스크립트는 향후 미세 조정 (fine-tuning)을 위해 RedPajama 데이터의 일부를 사전 토큰화 (pre-tokenize)하는 데 사용됩니다.

TARGET_MODEL=meta-llama/Llama-2-7b-hf # 토큰화에 사용됨
SEQLEN=4096
DATASET=togethercomputer/RedPajama-Data-1T-Sample
...

토큰화된 데이터셋은 특정 모델 제품군(또는 더 구체적으로는 해당 모델의 토크나이저 (tokenizer))에 특화되어 있습니다. 예를 들어, Llama-3 8B는 Llama-3 70B와는 호환되지만, 다른 토크나이저를 사용하기 때문에 Llama-2와는 호환되지 않습니다. 다른 모델을 위해 데이터를 토큰화하려면 1) --base_model, 2) model_seqlen, 3) --save_dataset_and_exit 경로를 설정해야 합니다.

또한 --preprocessing_num_workers를 하드웨어에 적합한 값으로 설정할 수 있습니다. --download_num_workers > 1로 설정하면 속도 제한 (rate limit) 등으로 인해 다운로드 오류가 발생할 수 있음에 유의하세요. 이러한 파라미터 및 기타 파라미터들은 스크립트의 --help에서 설명됩니다. 이 작업은 데이터셋 샘플과 전처리 캐시 (preprocessing cache)를 저장하기 위해 150-200 GiB의 디스크 공간을 필요로 합니다. 두 가지 모두 ./cache_dir에 저장되며 작업 후 삭제할 수 있습니다.

참고: 이전 미세 조정 (21년 8월 이전) 방식으로 결과를 재현하려면 커밋(commit) 559a366을 사용하세요.
이전 버전의 미세 조정은 PV-tuning 없이도 새로운 버전보다 결과는 좋지 않았지만, 속도는 더 빨랐습니다.

양자화된 모델의 정확도는 미세 조정 (fine-tuning)을 통해 더욱 향상될 수 있습니다.

우리의 새로운 PV-Tuning 알고리즘을 사용하려면, 스크립트를 실행하는 명령어가 다음과 같아야 합니다:

torchrun --nproc-per-node=$NUM_GPUS finetune.py \
--base_model $MODEL_PATH \
--quantized_model $QUANTIZED_WEIGHTS_PATH \
...

제로샷 평가(zero-shot evaluation)를 수행하기 위해, 우리는 Language Model Evaluation Harness 프레임워크를 채택합니다. 우리의 코드는 표준 transformers 형식의 모델과 호환되며, --aqlm_checkpoint_path 인자를 통해 (선택적으로) 양자화된 모델의 가중치를 로드할 수 있습니다.

PV-Tuning에서의 평가 결과는 lm-eval=0.4.0을 사용하여 생성되었습니다.

평가를 실행하려면 적절한 버전이 설치되어 있는지 확인하거나, 다음 명령어를 통해 설치하십시오:
pip install lm-eval==0.4.0

평가 절차를 시작하기 위한 메인 스크립트는 lmeval.py입니다.

export CUDA_VISIBLE_DEVICES=0,1,2,3 # 선택 사항: GPU 선택
export QUANTIZED_MODEL=<MAIN.py에서 저장된 양자화된 모델의 경로>
export MODEL_PATH=<허브에 있는 원본 모델의 경로를 입력하세요>
...

모델을 Hugging Face 호환 형식으로 변환하려면, 다음의 해당 인자들과 함께 convert_to_hf.py model in_path out_path를 사용하십시오:

model

• 원본 사전 학습된 모델 (main.py의 MODEL_PATH에 해당하며, 예: meta-llama/Llama-2-7b-hf). in_path