[논문 리뷰] Flashlight: Enabling Innovation in Tools for Machine Learning

Jacob Kahn, Vineel Pratap외 12명|arXiv (Cornell University)|2022. 01. 29.

Explainable Artificial Intelligence (XAI)인용 수 14

한 줄 요약

Flashlight는 새로운 텐서 연산, 메모리 관리 및 분산 학습 기법의 빠른 프로토타이핑을 가능하게 하여 시스템 연구를 가속화하기 위해 설계된 경량이며 모듈식인 머신러닝 라이브러리입니다. 최소한의 오버헤드와 내부 요소에 대한 완전한 제어를 통해 연구자들이 텐서 덧셈과 같은 핵심 연산을 사용자 정의 구현으로 쉽게 교체할 수 있으며, PyTorch 및 TensorFlow와 경쟁 가능한 성능을 달성하면서도 메인스트림 프레임워크에 기여할 수 있도록 합니다.

ABSTRACT

As the computational requirements for machine learning systems and the size and complexity of machine learning frameworks increases, essential framework innovation has become challenging. While computational needs have driven recent compiler, networking, and hardware advancements, utilization of those advancements by machine learning tools is occurring at a slower pace. This is in part due to the difficulties involved in prototyping new computational paradigms with existing frameworks. Large frameworks prioritize machine learning researchers and practitioners as end users and pay comparatively little attention to systems researchers who can push frameworks forward -- we argue that both are equally important stakeholders. We introduce Flashlight, an open-source library built to spur innovation in machine learning tools and systems by prioritizing open, modular, customizable internals and state-of-the-art, research-ready models and training setups across a variety of domains. Flashlight allows systems researchers to rapidly prototype and experiment with novel ideas in machine learning computation and has low overhead, competing with and often outperforming other popular machine learning frameworks. We see Flashlight as a tool enabling research that can benefit widely used libraries downstream and bring machine learning and systems researchers closer together. Flashlight is available at https://github.com/flashlight/flashlight .

연구 동기 및 목표

• 프레임워크의 크기와 복잡성이 점점 커짐에 따라 혁신하기 어려워지는 문제를 해결하기 위해.
• 머신러닝 연구자와 시스템 연구자 간 격차를 해소하기 위해 시스템 혁신에 최적화된 프레임워크를 제공하기 위해.
• 사용자 정의 메모리 관리자와 분산 계산 모델을 포함한 새로운 계산 패러다임의 신속한 프로토타이핑을 가능하게 하기 위해.
• PyTorch 및 TensorFlow와 같은 단일 프레임워크에 대체할 수 있는 최소화되고 모듈식인 대안을 제공하여 내부 컴포onent에 대한 종단 간 실험을 가능하게 하기 위해.
• 생산 환경에서 테스트되고 성능이 우수한 환경에서 신기술를 검증함으로써 메인스트림 프레임워크에의 기여를 촉진하기 위해.

제안 방법

• 텐서 연산과 계산 그래프를 위한 깔끔하고 조합 가능한 API를 갖춘 최소화된 모듈식 C++ 라이브러리로서 Flashlight를 설계하기 위해.
• 다양한 분야(예: 자연어 처리, 컴퓨터 비전)에서 완전하고 생산용 참조 모델을 구현하여 시스템 연구의 실험 기반을 마련하기 위해.
• 플러그인 방식의 메모리 관리자와 사용자 정의 할당기 백엔드를 통해 메모리 할당 및 텐서 레이아웃에 대한 저수준 제어를 제공하기 위해.
• 다양한 장치 매핑과 계산 스케줄링을 지원하여 새로운 분산 학습 전략에 대한 실험을 가능하게 하기 위해.
• 기존 연산자나 벤치마크를 수정하지 않고도 상속 또는 조합을 통해 핵심 연산(예: 텐서 덧셈)을 간편하게 교체할 수 있도록 하기 위해.
• 실험 중에 연산 수준에서 메모리 동작과 성능을 모니터링하기 위해 사용자 정의 텔레메트리 및 프로파일링 도구를 통합하기 위해.

실험 결과

연구 질문

• RQ1 최소화되고 모듈식인 머신러닝 라이브러리가 단일 프레임워크에 비해 시스템 수준 혁신을 얼마나 잘 지원할 수 있는가?
• RQ2 할당 및 분할에 대한 완전한 제어를 통해 연구자들이 새로운 메모리 관리 기법을 얼마나 효과적으로 프로토타이핑하고 벤치마크할 수 있는가?
• RQ3 사용자 정의 텐서 연산(예: 최적화된 덧셈)을 기존 코드를 손상시키지 않고 프레임워크에 얼마나 잘 통합할 수 있는가?
• RQ4 일반화된 분산 계산 및 옵티마이저 상태 스파이싱 접근 방식을 얼마나 민첩하고 확장 가능한 프레임워크에서 실제로 구현하고 평가할 수 있는가?
• RQ5 큰 복잡한 코드베이스에 의존하지 않고도 새로운 계산 원자요소에 대한 연구를 종단 간으로 대규모로 테스트할 수 있는 방법은 무엇인가?

주요 결과

• Flashlight는 PyTorch 및 TensorFlow에 비해 코드베이스와 바이너리 크기가 수 개의 주요 단위로 줄어들어 반복 속도가 빠르고 유지보수가 용이합니다.
• 연구자들은 큰 캐시 블록 분할을 제한하는 메모리 관리자를 신속하게 프로토타이핑하고 테스트하여 여러 모델에서 내부 분할을 20퍼센트 이상 감소시켰습니다.
• 사용자 정의 텐서 연산—예를 들어 새로운 텐서 덧셈의 구현—은 기존 코드를 수정하지 않고도 모든 연산자에서 원활하게 교체 가능하여 완전한 상호운용성을 입증했습니다.
• 프레임워크의 모듈식 설계 덕분에 버퍼 스파이싱 및 동적 장치 스케줄링를 포함한 일반화된 분산 학습 기법 개발이 가능했으며, 하드코딩된 가정이 전혀 없었습니다.
• 성능 벤치마크 결과, Flashlight의 참조 구현은 종종 PyTorch 및 TensorFlow의 엔드 투 엔드 학습 워크로드에서 성능가능하거나 이를 능가했습니다.
• Flashlight 연구에서 도출된 결과는 PyTorch에 성공적으로 기여되었으며, 캐싱 할당기 개선 사항이 포함되어 프레임워크의 실제 영향력을 입증했습니다.