[논문 리뷰] Automap: Towards Ergonomic Automated Parallelism for ML Models

Michael Schaarschmidt, Dominik Grewe외 9명|arXiv (Cornell University)|2021. 12. 06.

Parallel Computing and Optimization Techniques인용 수 5

한 줄 요약

Automap는 인도적 편의성과 컴파일러 통합된 SPMD 스타일 병렬 처리를 가능하게 하는 프로토타입 자동 분할기로, 인도적 편의성, 검색, 학습을 융합하여 트랜스포머 모델의 전문가 수준의 스플리팅 전략(예: 메가트론의 전략)을 복원한다. 플랫폼에 종속되지 않는 IR과 점진적 재작성 기반으로, 몇 분 내로 전문가 수준의 성능을 달성하며 수동 작업을 줄이고 복잡한 모델에 최소한의 사용자 입력으로도 확장 가능하다.

ABSTRACT

The rapid rise in demand for training large neural network architectures has brought into focus the need for partitioning strategies, for example by using data, model, or pipeline parallelism. Implementing these methods is increasingly supported through program primitives, but identifying efficient partitioning strategies requires expensive experimentation and expertise. We present the prototype of an automated partitioner that seamlessly integrates into existing compilers and existing user workflows. Our partitioner enables SPMD-style parallelism that encompasses data parallelism and parameter/activation sharding. Through a combination of inductive tactics and search in a platform-independent partitioning IR, automap can recover expert partitioning strategies such as Megatron sharding for transformer layers.

연구 동기 및 목표

• 대규모 머신러닝 모델에 대해 전문 지식이나 수동 재작성 없이 효율적인 데이터 병렬, 모델 병렬, 파이프라인 병렬 전략을 자동으로 발견하는 것.
• XLA와 같은 기존 컴파일러 스택과 JAX 워크플로우에 원활하게 통합되어 특정 하드웨어에 대해 재구현이 필요 없도록 하는 것.
• 최소한의 사용자 입력으로 빠른 상호작용 기반 분할을 가능하게 하여 실험 시간과 비용을 줄이고, 빠른 연구 반복을 지원하는 것.
• 반복 블록 등 구조적 특성을 가진 복잡하고 비정형적인 모델(수천 개 이상의 연산 포함)에 대해, 인도적 편의성과 모델 구조를 활용해 확장 가능한 분할 전략을 도출하는 것.
• 구조적 재작성 없이도 다양한 가속기 풍경(구형 하드웨어 포함)에서 효율적인 분할 전략을 지원함으로써 실질적인 배포를 가능하게 하는 것.

제안 방법

• 시스템은 XLA HLO 위에 쌓인 플랫폼에 종속되지 않는 새로운 분할 IR(MLIR 내의 다이어렉트)를 사용하여 스플리팅 및 복제 결정을 재작성 규칙으로 표현한다.
• 몬테카를로 트리 검색(MCTS)과 학습된 모델을 융합한 하이브리드 접근 방식을 적용하여, 고성능 전략을 찾기 위해 필요한 시도 횟수를 줄인다.
• 전문가 스플리팅 패턴(예: 메가트론)을 모방함으로써 인도적 편의성을 통합하여 검색 공간을 제약하고 수렴 속도를 향상시킨다.
• 점진적 재작성 기능을 지원하여, 특히 공유 상수와 레이어 그룹을 통해 계산 그래프 전반에 걸쳐 스플리팅 결정을 전파한다.
• 사용자는 'attention-block'과 같은 이름이 지정된 스코프를 통해 고수준 힌트를 제공함으로써 반복되는 레이어를 그룹화할 수 있으며, 이는 검색 효율성과 신뢰성을 크게 향상시킨다.
• 프로토타입은 JAX에서 인터랙션 네트워크를 사용해 합성 데이터로 훈련되었으며, Optax를 통한 최적화와 그래프 신경망 추론을 활용해 분할 결정을 내린다.

실험 결과

연구 질문

• RQ1 최소한의 사용자 입력으로도 메가트론과 같은 전문가 최적화 전략 수준의 성능을 유지하면서 자동 분할을 달성할 수 있는가?
• RQ2 수천 개 이상의 연산을 포함한 큰 비정형 모델에서 검색 효율성을 어떻게 향상시킬 수 있는가?
• RQ3 전문가 실천에서 유도된 인도적 편의성이 검색 공간을 얼마나 줄이고 수렴 속도를 얼마나 가속화하는가?
• RQ4 컴파일러 힌트(예: 이름이 지정된 레이어 그룹)는 공유 상수 간의 취약한 종속성 전파를 대체하여 깊은 네트워크에서 신뢰성을 향상시킬 수 있는가?
• RQ5 학습과 검색의 조합이 실제 연구 워크플로우에서 솔루션 도달 시간에 어떤 영향을 미치는가?

주요 결과

• 시스템은 500회 이내의 검색 에피소드(약 수 분) 내로 메가트론 스타일의 스플리팅을 발견하여, TPU v3에서 메가트론 수준의 성능을 거의 달성했으며, 런타임 오버헤드는 소량에 그쳤다.
• MCTS를 지도하는 데 학습된 필터를 사용함으로써, 근사 최적해를 안정적으로 찾기 위해 필요한 에피소드 수를 줄였으며, 데이터 기반 검색의 가치를 입증했다.
• 컴파일러 힌트(예: 이름이 지정된 레이어 그룹)를 사용하면, 공유 상수 간 전파에 의존하지 않더라도 메가트론 스플리팅을 소수의 에피소드 내로 안정적으로 발견할 수 있었다.
• 그룹화나 전파 기능 없이 24층의 트랜스포머에서 메가트론 스플리팅을 발견하지 못했으며, 이는 구조적 힌트가 확장성에서 핵심적인 역할을 한다는 점을 강조한다.
• 기존의 표준 모델이 아닌 모델(예: GraphNets)도 성공적으로 분할하여, 더 큰 규모의 실험을 위한 입력 엣지 스플리팅과 같은 실용적인 전략을 발견했다.
• 5만~10만 개의 연산을 포함하는 모델에 대해서도 이론적으로 확장 가능하며, 상호작용 연구 워크플로우에 적합한 검색 시간(수 분, 수 시간이 아님)을 확보했다.