HomebrewML/HeavyBall

해결하는 문제

HeavyBall은 PyTorch를 위한 고성능의 조합 가능한 최적화 라이브러리를 제공합니다. 컴파일된 빌딩 블록을 사용하여 Triton 커널로 융합함으로써 메모리 트래픽을 줄이고 실행 속도를 높여 표준 최적화 구현의 비효율성을 해결합니다. 또한 구현하거나 조합하기 어려운 고급 최적화 기법(예: 2차 방법 및 상태 압축)의 통합을 간소화합니다.

작동 방식

이 라이브러리는 100개 이상의 컴파일된 함수들로 구성되며, 이들은 변환 체인으로 조립됩니다. torch.compile(fullgraph=True)를 사용하면 이러한 함수들이 최소한의 커널로 융합되어 메모리에 대한 읽기 및 쓰기 횟수를 크게 줄입니다. 예를 들어, 표준 Adam 업데이트는 14번의 읽기와 9번의 쓰기에서 4번의 읽기와 3번의 쓰기로 줄일 수 있습니다.

1차(AdamW, SGD), 직교(Muon), Shampoo 기반(SOAP) 및 크로네커 분해(PSGD)를 포함한 광범위한 유형의 최적화기를 지원합니다. 또한 서로 다른 매개변수 그룹에 서로 다른 최적화기를 적용할 수 있는 "SplitOpt" 기능도 제공합니다. 분산 학습의 경우, FSDP를 사용할 때 2차 방법에 대한 재분할을 자동으로 처리합니다.

대상 사용자

PyTorch 생태계 내에서 더 빠른 최적화 단계, 최적화 상태에 대한 더 낮은 메모리 오버헤드, 또는 다양한 고급 2차 및 직교 최적화 알고리즘에 대한 액세스가 필요한 머신러닝 엔지니어 및 연구자.

주요 기능

• 광범위한 최적화기 제품군: AdamW, SGD 및 RMSprop에 대한 API 호환 대체제와 함께 Muon, SOAP, LATHER 및 ADOPT를 포함합니다.
• 조합 가능한 기능: MARS 분산 감소, 신중한 업데이트 및 PaLM 스타일 beta2 스케줄링을 위한 체인 가능한 플래그입니다.
• ECC 상태 압축: 정확도를 희생하지 않고 공간을 절약하기 위해 최적화 상태 메모리 사용량(예: bf16 + int8 보정)을 줄입니다.
• 고성능: torch.compile을 통해 연산을 Triton 커널로 융합하여 단계 지연 시간을 크게 단축합니다.
• 분산 지원: 복잡한 2차 방법에 대한 자동 재분할을 포함하여 DDP 및 FSDP와의 기본 호환성을 제공합니다.