오픈 소스 커뮤니티를 위한 AI 모델 훈련을 개선하는 IBM® Research 데이터 로더

대량의 데이터에 대해 AI 모델을 학습할 때 병목 현상을 어떻게 극복할 수 있을까요? 올해의 PyTorch 컨퍼런스에서 IBM Research는 대규모 LLM 교육을 위한 획기적인 데이터 로더를 선보였습니다. 이제 PyTorch 사용자가 사용할 수 있는 이 도구는 가능한 한 많은 사용자를 대상으로 대규모 교육을 간소화하는 것을 목표로 합니다.

고처리량 데이터 로더에 대한 아이디어는 연구 과학자들이 모델 학습 중에 관찰한 실질적인 문제에서 비롯되었습니다. 그들의 작업에는 점점 더 효율적인 GPU에 맞춰가면서 여러 장치에서 대량의 데이터를 처리할 수 있는 도구가 필요했습니다. IBM Research는 블로그에서 이번 릴리스에 대해 "단순히 작업을 완료하는 데 필요한 도구를 구축한 연구팀 덕분"이라고 언급했습니다.

IBM Research의 Davis Wertheimer는 대규모 교육 시 발생할 수 있는 몇 가지 문제를 다음과 같이 설명합니다. "대규모 교육에는 80/20 규칙이 있습니다. 발표된 모든 문헌의 80%는 GPU 메모리와 통신 및 연산 간의 알고리즘 트레이드오프에 대해 살펴보고 있습니다. 하지만 실제로 무언가를 구축하려고 할 때, 80%, 파이프라인이 가장 좁은 병목현상의 속도로 실행되기 때문에 다른 모든 실질적인 문제들의 긴 꼬리에 의존할 수 있습니다."

IBM 팀은 교육 플랫폼을 개발하면서 계속해서 병목 현상을 겪었습니다. "GPU 사용이 점점 더 능숙해짐에 따라 데이터 로더에서 병목 현상이 발생하는 경우가 점점 더 많아지고 있습니다."라고 Wertheimer는 말합니다.

이러한 깨달음은 이중 개발 프로세스로 이어졌습니다. "교육 플랫폼을 발전시키는 한편, 병목 현상을 방지하기 위해 교육 플랫폼의 속도 요구 사항을 따라잡기 위해 데이터 로더를 지속적으로 발전시키는 과정이 병행되어 왔습니다."라고 그는 설명합니다.

IBM Research의 Linsong Chu는 데이터 로더의 필수 기능을 다음과 같이 간략하게 설명합니다.

상태 저장 및 체크포인트 가능: "모델을 저장할 때마다 데이터 로더 상태도 저장되며, 체크포인트에서 복구할 때마다 모델 상태와 데이터 로더 상태를 동시에 복구해야 합니다."라고 Chu는 말합니다.

체크포인트 자동 재조정: 데이터 로더는 확장된 교육 세션 중에 워크로드 변화에 맞게 자동으로 조정됩니다. "교육에는 몇 주 또는 몇 달이 쉽게 걸릴 수 있으며, 중간에 워크로드의 규모를 조정해야 하는 데에는 수많은 이유가 있습니다."라고 Chu는 말합니다.

데이터 스트리밍: 시스템은 셔플링에 대한 빌드 오버헤드가 전혀 없는 데이터 스트리밍을 지원합니다.

비동기식 분산 작업: "우리는 데이터 로더가 비차단 기능이기를 원합니다."라고 Chu는 설명합니다. "데이터 로더 상태를 저장하는 동시에 통신이 전혀 필요하지 않은 형태로 저장이 분산되기를 원합니다."

동적 데이터 혼합: 데이터 로더는 다양한 데이터 혼합 비율에 적응할 수 있으므로 진화하는 학습 요구 사항에 유용합니다.

효율적인 글로벌 셔플: 이 도구는 대규모 데이터 세트를 처리할 때 메모리 병목 현상을 해결하여 데이터가 증가하더라도 셔플을 효율적으로 만듭니다.

PyTorch 네이티브, 모듈식 및 확장성: 적응성과 확장성을 위해 설계된 데이터 로더는 향후 성장에 대비할 수 있습니다. "내년에 30조, 50조 또는 100조 토큰을 처리해야 한다면 어떨까요?"라고 Chu는 묻습니다. "세상은 빠르게 변화하고 있습니다. 따라서 우리는 데이터 로더가 오늘뿐만 아니라 내일을 위해 살아남을 수 있도록 구축해야 합니다."

IBM Research 팀은 수백 개의 크고 작은 작업을 실행하면서 몇 달에 걸쳐 데이터 로더를 엄격하게 테스트했습니다. 그들은 안정적이고 원활한 코드 번호를 관찰했습니다. 게다가 전체 데이터 로더는 비동기적으로 작동하며 차단되지 않습니다.

"이 모든 것을 실현하기 위해 내장된 PyTorch 기능을 많이 활용했습니다."라고 Wertheimer는 말합니다. "이것이 바로 우리가 다시 기여하는 이유입니다."