에
v0.6.1
소개 | 추론 | 훈련 | 구성성 | 사용자 정의 커널 | 알파 기능 | 설치 | 통합 | 비디오 | 라이센스 | 소환
torchao: 사용자 정의 데이터 유형 및 최적화를 위한 PyTorch 라이브러리입니다. 추론과 훈련을 위해 가중치, 기울기, 최적화 도구 및 활성화를 양자화하고 희소화합니다.
빠른 시리즈를 가져온 팀에서
torchao는 Huggingface의 대부분의 PyTorch 모델에 대해 torch.compile() 및 FSDP2 와 함께 작동합니다.
모델을 양자화하고 스파스화하는 것은 좋아하는 HuggingFace 모델을 포함하여 nn.Linear 가 있는 모든 모델에서 작동해야 하는 1개의 라이너입니다. 보다 포괄적인 사용 지침은 여기에서, 희소성은 여기에서, HuggingFace 추론 예는 여기에서 찾을 수 있습니다.
추론을 위해 다음과 같은 옵션이 있습니다.
from torchao . quantization . quant_api import (
quantize_ ,
int8_dynamic_activation_int8_weight ,
int4_weight_only ,
int8_weight_only
)
quantize_ ( m , int4_weight_only ()) gpt-fast의 경우 int4_weight_only() tok/s가 2배이고 VRAM 요구 사항이 torch.compiled 기준에 비해 약 65% 줄어들기 때문에 bs=1에서 가장 좋은 옵션입니다.
GPU에서 전체 모델을 양자화할 만큼 VRAM이 충분하지 않고 CPU 양자화가 너무 느린 경우 quantize_(model, int8_weight_only(), device="cuda") 와 같은 장치 인수를 사용할 수 있습니다. 각 레이어를 GPU에 개별적으로 양자화합니다.
이러한 기술 중 속도가 느려지는 경우 또는 어떤 옵션을 사용해야 할지 확실하지 않은 경우 자동으로 레이어를 프로파일링하고 각 레이어를 양자화하는 가장 좋은 방법을 선택하는 자동퀀트 사용을 고려해 보세요.
model = torchao . autoquant ( torch . compile ( model , mode = 'max-autotune' ))우리는 또한 개발자용 API를 제공하므로 귀하가 자신의 양자화 알고리즘을 구현할 수 있으므로 우수한 HQQ 알고리즘을 동기 부여 사례로 사용하십시오.
긴 컨텍스트 길이(그리고 필연적으로 메모리 효율성) 추론을 가능하게 하기 위해 kv 캐시 양자화 및 기타 기능을 추가했습니다.
실제로 int4 가중치 전용 양자화와 함께 이러한 기능을 사용하면 최대 메모리를 ~55%까지 줄일 수 있습니다. 즉, 최대 메모리 18.9GB만으로 130k 컨텍스트 길이로 Llama3.1-8B 추론을 수행할 수 있습니다. 자세한 내용은 여기에서 확인할 수 있습니다.
훈련 후 양자화는 빠르고 컴팩트한 모델을 만들 수 있지만 정확도가 저하될 수도 있습니다. 이러한 제한을 극복하려면 QAT(Quantization Aware Training)를 살펴보는 것이 좋습니다. Torchtune과 협력하여 우리는 훈련 후 양자화(PTQ)와 비교하여 Llama3에 대한 hellaswag의 정확도 저하 96%와 wikitext의 난관 저하 68%를 복구하여 기존 PTQ에 비해 상당한 정확도 향상을 보여주는 QAT 레시피를 개발했습니다. . 여기에 전체 레시피를 제공했습니다.
from torchao . quantization . qat import Int8DynActInt4WeightQATQuantizer
qat_quantizer = Int8DynActInt4WeightQATQuantizer ()
# Insert "fake quantize" operations into linear layers.
# These operations simulate quantization numerics
model = qat_quantizer . prepare ( model )
# Run Training...
# Convert fake quantize to actual quantize operations
model = qat_quantizer . convert ( model )torchao.float8은 https://arxiv.org/abs/2209.05433에 설명된 대로 확장된 float8 dtype을 사용하여 훈련 레시피를 구현합니다.
torch.compile 켜면 현재 결과에 따르면 128개의 H100 GPU LLaMa 3 70B 사전 훈련 작업에서 처리량 속도가 최대 1.5배 향상되었습니다(세부 정보).
from torchao . float8 import convert_to_float8_training
convert_to_float8_training ( m , module_filter_fn = ...)그리고 float8을 사용한 사전 훈련의 최소 수준의 훈련 방법에 대해서는 torchtitan을 확인하세요.
ViT-L에서 6%의 종단 간 속도 향상 으로 반구조화된 2:4 희소성에 대한 지원을 추가했습니다. 전체 블로그는 여기
코드 변경은 여기에서 전체 예제를 볼 수 있는 1개의 라이너입니다.
swap_linear_with_semi_sparse_linear ( model , { "seq.0" : SemiSparseLinear })ADAM은 모델 매개변수보다 2배 많은 메모리를 사용하므로 최적화 기능 상태를 8비트 또는 4비트로 양자화하여 최적화 기능 VRAM 요구 사항을 fp16 기준에 비해 각각 2배 또는 4배 줄일 수 있습니다.
from torchao . prototype . low_bit_optim import AdamW8bit , AdamW4bit , AdamWFp8
optim = AdamW8bit ( model . parameters ()) # replace with Adam4bit and AdamFp8 for the 4 / fp8 versions실제로 우리는 전문적으로 작성된 커널보다 약간 느리지만 이러한 최적화 프로그램의 구현은 수백 줄의 PyTorch 코드 로 작성되고 컴파일되었으므로 양자화된 최적화 프로그램에 사용하거나 복사하여 붙여넣으십시오. 벤치마크는 여기
또한 그라디언트(가중치와 동일한 크기)와 최적화 프로그램이 모두 효율적으로 CPU로 전송되는 단일 GPU CPU 오프로딩을 지원합니다. 이것만으로도 VRAM 요구 사항을 60%까지 줄일 수 있습니다.
optim = CPUOffloadOptimizer ( model . parameters (), torch . optim . AdamW , fused = True )
optim . load_state_dict ( ckpt [ "optim" ])torch.compile : 우리의 핵심 디자인 원칙은 우리가 제공하는 새로운 dtype이나 레이아웃에서 컴파일러와 함께 작동해야 하는 구성 가능성입니다. 커널이 순수 PyTorch, CUDA, C++ 또는 Triton으로 작성되었는지 여부는 중요하지 않습니다. 모든 것이 제대로 작동해야 합니다! 그래서 우리는 순수한 PyTorch에서 dtype, 레이아웃 또는 비트 패킹 논리를 작성하고 효율적인 커널을 코드 생성합니다.낮은 비트 dtype의 구성 가능성을 compile 및 fsdp와 결합한 가장 좋은 예는 QLoRA 알고리즘을 구현하는 데 사용한 NF4입니다. 따라서 귀하가 이 분야의 교차점에서 연구를 수행 중이라면 우리는 귀하의 의견을 듣고 싶습니다.
우리는 torch.compile() 사용하여 그래프 중단이 없는 사용자 정의 작업을 작성하고 릴리스하기 위한 지원을 추가했습니다. 따라서 커널 작성을 좋아하지만 커널을 패키징하여 모든 운영 체제와 Cuda 버전에서 작동하는 것을 싫어한다면 다음에 대한 기여를 환영합니다. 귀하의 맞춤 작전. 따라할 수 있는 몇 가지 예가 있습니다.
quantize_(model, fpx_weight_only(3, 2))우리가 자세히 살펴봐야 할 다른 CUDA 커널이 있다고 생각한다면 이 문제에 대한 의견을 남겨주세요.
우리가 기대하고 있지만 오븐에서 요리하는 데 시간이 더 필요한 것들
quantize_(model, int8_weight_only_quantized_training()) 과 함께 사용하기 쉽습니다. 메모리 벤치마크가 아직 설득력이 없기 때문에 이 작업은 프로토타입입니다. torchao Pytorch의 몇 가지 새로운 기능을 자유롭게 사용하므로 현재 Nightly 또는 최신 안정 버전의 PyTorch와 함께 사용하는 것이 좋습니다.
기본적으로 CUDA 12.1로 설정되는 Pypi의 안정적인 릴리스
pip install torchaoPyTorch 인덱스의 안정적인 릴리스
pip install torchao --extra-index-url https://download.pytorch.org/whl/cu121 # full options are cpu/cu118/cu121/cu124야간 출시
pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu121 # full options are cpu/cu118/cu121/cu124대부분의 개발자는 더 빠른 반복을 위해 사용자 정의 C++/CUDA 확장 빌드를 건너뛰고 싶을 것입니다.
USE_CPP=0 pip install -e . 우리는 또한 다음을 포함하는 일부 주요 오픈 소스 라이브러리에 통합될 수 있게 된 것을 행운으로 생각합니다.
torchao BSD 3 라이센스로 출시되었습니다.
Torchao 라이브러리가 유용하다고 생각되면 아래와 같이 작업에 인용해 주세요.
@software { torchao ,
title = { torchao: PyTorch native quantization and sparsity for training and inference } ,
author = { torchao maintainers and contributors } ,
url = { https//github.com/pytorch/torchao } ,
license = { BSD-3-Clause } ,
month = oct,
year = { 2024 }
