컴퓨팅을 위한 데이터 스케줄링: AMD GPU 기반 LLM 학습을 위한 계층적 데이터 파이프라이닝

대규모 언어 모델(LLM)은 최근 머신러닝 혁신의 핵심 초점입니다. 모델 크기가 기하급수적으로 증가함에 따라, 학습에 필요한 계산량은 천문학적 수준입니다. GPU는 막대한 양의 산술 연산을 수행할 수 있습니다. 하지만 그 컴퓨팅 유닛이 필요로 하는 데이터가 항상 준비되는 것은 아닙니다. 이 데이터는 여전히 느린 메모리나 다른 GPU에 있을 수 있습니다. 칩이 기다리는 동안, 값비싼 컴퓨팅 자원이 유휴 상태로 남아있게 되며, 대규모 환경에서는 활용률의 작은 하락조차도 상당한 비용으로 이어집니다. 본 논문에서 우리는 이러한 대기 시간을 주요 병목 현상으로 다룹니다. 우리는 명시적인 데이터 이동을 위해 설계된 하드웨어 기능에 초점을 맞춥니다. 이 기능들은 칩 외부 메모리(off-chip memory)에서 데이터를 칩 내부로 로드하고, 로드되는 과정에서 컴퓨팅이 요구하는 레이아웃으로 재배열하며, GPU 간에 데이터를 전송합니다. 이러한 기능들을 오케스트레이션(조율)함으로써, 우리는 데이터가 컴퓨팅 유닛으로 꾸준히 흐르는 파이프라인을 유지하여 유닛들이 바쁘게 작동하도록 합니다. 우리는 세 가지 핵심 데이터 이동 명령어와 하드웨어 기능을 사용하여 이 병목 현상을 해결하는 커널 및 프레임워크 구현을 제시합니다. Tensor Data Mover (TDM)는 고대역폭 메모리(HBM)에서 칩 내부 메모리로 텐서 데이터를 이동시킵니다. Transpose-load 명령어는 LLM 컴퓨팅이 요구하는 레이아웃으로 텐서 데이터를 로드되는 과정에서 전치(transpose)합니다. 마지막으로, System DMA (SDMA)와 Remote DMA (RDMA)는 각각 스케일업 및 스케일아웃 네트워크를 통해 GPU 간에 데이터를 이동시킵니다. 우리의 결과는 TDM과 Transpose-load가 커널 수준에서 평균 1.88배 및 1.19배의 속도 향상을 제공함을 보여줍니다. SDMA와 RDMA는 통신과 중첩된 컴퓨팅에서 각각 평균 1.16배 및 1.33배의 속도 향상을 가져옵니다. 이는 모든 GPU 데이터 파이프라이닝 기능을 활용하는 최초의 AMD LLM 학습 시스템입니다. 우리의 구현은 오픈 소스 프레임워크 및 라이브러리에 통합되었습니다. https://github.com/AMD-AGI/Primus-Turbo https://github.com/AMD-AGI/Primus