pytorch/ao

해결하는 문제

TorchAO는 훈련에서 서비스까지 모델 최적화를 위한 네이티브 PyTorch 라이브러리를 제공합니다. 모델 크기를 줄이고 추론 속도를 높이되, 정확도를 크게 훼손하지 않도록 하는 도전 과제를 해결하며, 특히 양자화와 희소성에 대한 도구를 제공합니다.

작동 방식

TorchAO는 torch.compile() 및 FSDP2에 직접 통합되어 대부분의 HuggingFace 모델에서 즉시 작동합니다. PyTorch, CUDA, C++, 또는 Triton의 커널을 사용하여 다양한 저정밀도 데이터 유형(dt)과 메모리 레이아웃을 구현하여 GPU 및 CPU에서 성능을 가속화합니다.

대상 사용자

대규모 모델(예: LLM 또는 디퓨전 모델)을 서버 GPU, ExecuTorch를 통한 엣지 디바이스, 또는 ARM CPU에 더 효율적으로 배포해야 하는 머신러닝 엔지니어 및 연구자들을 위한 것입니다.

주요 특징

• 양자화 인식 훈련(QAT): 훈련 후 양자화로 인한 정확도 저하를 복구하며, 특히 int4 모델에서 효과적입니다.
• Float8 훈련: 스케일된 float8 데이터 유형을 사용하여 사전 훈련(예: Llama-3.1-70B)을 최대 1.5배 빠르게 수행합니다.
• 추론 최적화: int4 가중치 전용 및 float8 동적 양자화를 지원하여 상당한 속도 향상과 메모리 절감을 실현합니다.
• 메모리 효율적인 옵티마이저: 양자화된 옵티마이저(AdamW 4비트/8비트/Fp8) 및 CPU 오프로딩을 포함하여 VRAM 요구량을 줄입니다.
• C 레벨 통합: vLLM, HuggingFace Transformers/Diffusers/PEFT, Unsloth, SGLang과 통합되어 있습니다.