GPU 클러스터의 영향력 있는 스케줄링

높은 영향력을 가진 연구를 우선순위로 지정하면서도 전체 점유율을 유지하는 클러스터 스케줄러 구축

Ai2의 AI 인프라 팀은 기관의 GPU 컴퓨팅 용량을 제공할 책임이 있으며, 특히 대규모 분산 훈련 워크로드에 초점을 맞추고 있습니다. 저희는 이 작업을 서로 쌓아 올리는 네 가지 지표의 피라미드로 생각합니다.

가장 기초는 **가용성(availability)**입니다. 하드웨어가 얼마나 자주 건강하고 작업 준비가 되어 있는지를 나타냅니다. 그 위에는 **점유율(occupancy)**이 있습니다. 이는 특정 워크로드에 할당된 사용 가능 시간의 비율을 의미합니다. 다음은 **영향력(impact)**으로, 가장 가치 있는 워크로드가 자원을 받기 위해 얼마나 자주 선택되는지를 나타냅니다. 피라미드의 정점은 **활용률(utilization)**이며, 이는 워크로드 수명 주기 동안 사용된 GPU 용량의 비율입니다.

이 게시물은 저희 스케줄링 결정의 영향력을 개선하는 것에 관한 것입니다. 최근에 저희는 우선순위 기반 스케줄러를 GPU 시간 예산, 계층적 공정 공유 할당(hierarchical fair-share allocation), 그리고 시간 슬라이싱 계약을 포함하는 시스템으로 교체했습니다. 그 결과, 각 연구 프로젝트가 얼마나 많은 GPU 시간을 받을 자격이 있는지에 대한 논의를 개별적인 운영 작업에서 투명한 관리 예산 책정 프로세스로 전환할 수 있었습니다.

Ai2에서는 88개부터 1024개까지 크기가 다양하고 NVIDIA H100, B200, B300 GPU가 클러스터 형태로 배치되어 관리되고 있습니다. 이 클러스터들은 AI 모델의 대규모 분산 훈련을 위해 구축되었으며, LLM 및 VLM 훈련의 전체 모델 흐름, 로보틱스 강화학습(RL) 시뮬레이션, 그리고 과학적 에이전트 사용 사례를 위한 사후 훈련을 포함하여 다양한 AI 도메인을 다루는 약 150명의 내부 연구원 그룹에 서비스를 제공합니다.

많은 연구실과 마찬가지로, 저희는 GPU 시간에 대한 수요가 공급을 훨씬 초과하는 상황입니다. 제출된 워크로드에 따르면, 어느 순간에도 사용 가능한 것보다 23배 더 많은 GPU에 대한 미처리 요청이 있습니다. 이를 생각하는 한 가지 방법은 저희 클러스터의 모든 사용 가능한 GPU 시간이 23개의 서로 다른 연구 워크로드가 경쟁하고 있다는 것입니다.

역사적으로 우리는 우선순위 기반 스케줄러를 사용했으며, 워크로드가 선점 가능성(preemptability)에서 제외되는 것을 허용했습니다. 각 팀은 선점으로부터 보호받는 워크로드에 의해 사용할 수 있는 동시 GPU 사용량 제한을 가지고 있었습니다. 선점 가능한 워크로드는 유휴 GPU에 대해 그 제한을 초과할 수 있었습니다. 이 전략은 예측 가능한 병리 현상을 낳았습니다. 예를 들어, 사용자들은 필요할 때 연결할 수 있도록 아무 작업도 수행하지 않는(no-op) 워크로드를 주차하는 'GPU 스쿼팅(squatting)' 사례를 관찰했습니다. 이는 연구자들이 실시간으로 문제를 해결하기에 충분히 낮은 지연 시간으로 디버깅 워크로드를 시작할 수 없다는 것을 발견했기 때문입니다. 또한, 결국 예약된 모든 워크로드의 100%가 높은 우선순위(HIGH priority)를 사용하는 '우선순위 인플레이션(priority inflation)'도 관찰되었습니다. 이는 낮은 우선순위 레벨이 GPU 시간을 완전히 박탈당한다는 것을 의미했습니다. 선점 가능성이 선택 사항이었기 때문에, 우리는 또한 당직 엔지니어들이 알려진 유지보수 문제가 있는 호스트에서 실행되는 비선점 워크로드의 조직적인 종료를 협상하는 데 대부분의 티켓 응답 시간을 소비한다는 사실을 발견했습니다.

이러한 문제들이 발생했을 때, 우리는 근본 원인을 파악하는 것이 느렸습니다. 가장 중요한 작업이 GPU 시간을 확보하도록 보장하려는 우리의 초기 시도는 우선순위가 설정되는 방식을 더 엄격하게 통제하는 데 초점을 맞추었으며, 궁극적으로는 중요한 프로젝트에 명시적으로 GPU 독점권을 할당하여 우선순위 기반 스케줄러를 우회하는 것이었습니다. 우리는 처음에는 인식하지 못했지만, '공유지의 비극(tragedy of the commons)'을 관찰하기 위한 완벽한 실험실을 구축했던 것입니다. 개인들은 희소하고 공유되는 자원을 놓고 경쟁했고, 개별적인 결과를 최대화하려고 노력함으로써 최적이 아닌 전역적인 결과에 도달했으며 근본적인 자원을 남용했습니다.

이러한 상호작용을 관찰한 것은 저희가 처음은 아니었습니다. 자원 할당(Resource allocation)은 알고리즘 개발, 경제학, 시스템 관리(system management)를 혼합하는 매력적인 연구 분야입니다. 핵심 문제는 사용자들이 자신의 작업 가치(value of their own jobs)를 조직보다 더 잘 아는 경우가 많지만, 그 가치를 숨기거나 자원을 붙잡아 두려는 인센티브가 있어 전체 성능에 손해를 입힐 수도 있다는 것입니다. 예를 들어, 2011년에 Dominant Resource Fairness를 소개한 Ghodsi 등의 연구진은 검색 회사가 사용자들로부터 높은 활용률(high utilization)을 보장받는 경우에만 전용 머신을 작업에 제공했던 일화를 언급합니다. 그들은 곧 “사용자들이 활용률 수준을 인위적으로 부풀리기 위해 코드에 무한 루프를 뿌리는 것”을 발견했습니다. 하드웨어는 바뀌었지만, 자원 할당을 복잡하게 만드는 근본적인 문제들은 계속 남아 있습니다.

공유지의 비극(tragedy of the commons)에 대한 고전적인 해결책은 공유 자원을 사유화하는 것입니다. 소유자들은 자신의 재산 가치를 극대화하도록 인센티브를 받습니다. 저희가 팀들에게 GPU 세트에 독점권(monopoly)을 할당했을 때, 이미 이러한 방식의 버전을 수행하고 있었지만, 너무 거칠었습니다. 이는 연구의 계절성 때문에 GPU들이 유휴 상태로 머물게 만들었습니다. 팀들은 서로 다른 시간에 실험과 훈련을 진행할 준비가 되어 있기 때문에, 독점권을 할당하면 어떤 작업도 실행될 준비가 되지 않아 자원을 기다리는 다른 팀이 생기는 시간이 있을 것이 분명했습니다.

저희는 정적인 스케줄에 동적으로 변화하는 연구 요구 사항들을 맞추려고 노력하며 배낭 문제(knapsack problem)를 수동으로 해결하고 있었습니다. 저희는 소유권 인센티브도 원했지만, GPU의 완전한 점유율(full occupancy)을 유지하는 것도 원했습니다.

우리는 소유권 모델을 개선하기로 결정했습니다. 팀에게 GPU를 할당하는 대신, 일정량의 GPU 시간을 배분하도록 선택했습니다. 미래 수요를 예측하려면 새로운 과학 실험의 결과를 알아야 하므로 정확하게 예측할 수 없습니다. 하지만 연구 노력 전반에 걸친 우선순위는 전략의 문제이며, 이는 사전에 더 쉽게 논의되고 결정될 수 있습니다. 스케줄링 퍼즐을 풀려고 하기보다, 리더십이 투자자처럼 생각하도록 만들었습니다. 워크로드가 존재하기 전에, 각 연구 노력이 가질 가능성이 있는 영향에 대한 판단을 바탕으로 GPU 시간을 어떻게 자금 지원할지 결정합니다. 그러면 스케줄러는 도착하는 워크로드의 우선순위를 정할 때 이 정보를 사용할 수 있습니다.

이 점을 염두에 두고, 관리자가 자신이 책임지는 프로젝트와 연구원들에게 GPU 시간을 비례적으로 할당할 수 있는 계층적 시스템을 고안했습니다. 아래 다이어그램에서 보여주듯이, 이는 프로그램 전략을 보장된 GPU 시간의 지분으로 직접 변환합니다. Project A1은 다른 곳에 대기 중인 다른 프로젝트가 얼마나 많든 상관없이 전체 용량의 35%를 확보한다는 것을 알고 있습니다.

*괄호 안 값은 리프(leaf) 프로젝트에 할당된 총 클러스터 용량을 나타냅니다.

이 시스템에서는 GPU 시간을 요청하는 모든 작업에 예산(budget)이 책정되어야 하며, 그렇지 않으면 선점(preemption)으로부터 보호받을 수 없습니다. 이전 시스템에서는 높은 우선순위(HIGH priority)가 비용이 들지 않았고 비선점성(non-preemptibility) 덕분에 팀이 자신들의 동시 GPU 한도를 무기한으로 채울 수 있었으므로, 모두가 이를 사용했습니다. 이제는 아무것도 공짜가 아니기 때문에, GPU 시간을 얻기 위한 모든 속임수는 이득을 보는 사용자 할당량에서 끌어옵니다. '자리 차지(squatting)' 워크로드는 아무것도 하지 않으면서 팀 예산을 소진하는 것을 의미합니다. 우리의 전략은 스케줄러를 이용해 게임을 하는 것이 더 큰 예산에 대해 정직하게 논의하는 것보다 비용이 많이 들도록 만드는 것입니다. 우리는 이 예산 검토 프로세스를 끊임없이 개선하고 있지만, 핵심 요구 사항은 연구원들이 필요한 시간을 주장할 수 있는 빈번한 기회가 있어야 하며, 결정은 문제의 트레이드오프(tradeoffs)에 대해 가장 많은 맥락을 가진 관리자들에 의해 이루어져야 한다는 것입니다. 이는 연구 프로젝트 내의 할당 결정은 책임 연구원(lead researcher)이, 연구 프로그램 내에서는 주임 연구원(principal investigator)이, 그리고 여러 프로그램을 아우르는 경우에는 책임 프로그램 관리자(lead program manager) 또는 CEO가 내린다는 것을 의미합니다.

이 GPU 시간 예산 도구와 결합하여, 우리는 프로그램 트리 전체에 걸쳐 실제 점유율을 관리하기 위한 계층적 공정 공유 스케줄러(hierarchical fair-share scheduler)를 구축했습니다. 여기서 알고리즘은 새로운 것이 아닙니다—시간 창(time window)에 대한 계층적 공정 공유는 2009년 Hadoop Fair Scheduler로 거슬러 올라가는 역사의 일부이며, 동일한 접근 방식이 오늘날 SLURM의 Fair Tree와 YARN의 Fair Scheduler에서 활발하게 사용되고 있습니다. 우리에게 새로운 것은 입력값입니다: 트리가 연구 프로그램 구조를 반영하며, 가중치(weights)는 정적 할당량(static quotas)이 아니라 관리자들이 설정하는 예산입니다.

스케줄러는 슬라이딩 룩백 창(sliding lookback window)(기본값은 7일)에 걸쳐 점유율을 추적하며, 활용도가 낮은 할당량의 워크로드를 활용도가 높은 할당량보다 먼저 정렬합니다. 이러한 방식으로, 일주일이라는 시간 범위 동안 모든 그룹이 충분한 수요를 가진 워크로드를 적극적으로 제출하는 한, 할당된 GPU 시간을 받을 것으로 예상할 수 있습니다.

새로운 스케줄러는 마치 컴퓨팅 자원이 30% 더 생긴 것처럼 느껴지게 합니다. 이전 스케줄러에서는 전체 슬롯 한도를 다 사용하지 않는 순간이 오면, 그 컴퓨팅 자원은 사실상 손실되었습니다. 하지만 새로운 스케줄러를 사용하면, 이런 상황이 발생했을 때 나중에 할당량 한도를 초과하여 버스트(burst)할 수 있고 여전히 작업들이 빠르게 그리고 선점(preemption) 없이 스케줄링되는 것을 볼 수 있습니다. 이는 본질적으로 우리가 그 컴퓨팅 자원을 되찾을 수 있게 해줍니다. 저희 워크로드는 종종 버스티(bursty)하기 때문에, 이것은 저희에게 상당한 양의 컴퓨팅 자원을 돌려주었습니다." — Chris Clark

스케줄러는 두 가지 종류의 점유율을 구분합니다. *할당된 점유율 (Allocated occupancy)*은 워크로드가 예산으로 청구되는 시간입니다. 이는 워크로드 소유자의 할당량에서 가져오며, 공정 공유(fair-share) 예산 계산에 영향을 미치고, 이러한 워크로드는 최소 실행 시간 창 동안 선점으로부터 보호됩니다. *할당되지 않은 점유율 (Unallocated occupancy)*은 어떤 예산에도 청구되지 않으며, 처음부터 보호받지 못하고 모든 할당된 요청에 의해 선점될 수 있습니다. 이를 통해 할당량이 수요와 제대로 일치하지 않을 때도 GPU를 완전히 점유 상태로 유지할 수 있게 하며, 팀들이 무료 GPU 사이클을 거절하는 것을 방지합니다.

분산 학습에서 공정한 자원 할당을 어렵게 만드는 추가적인 기능은 워크로드가 매우 오랫동안 실행될 수 있다는 것입니다. 학습 작업들은 정기적으로 몇 시간, 며칠, 때로는 몇 주 동안 실행됩니다. 일단 스케줄링되면, 하나의 워크로드는 지정된 GPU에 일주일 이상 남아 있을 수 있으며, 다른 사람들이 예산화된 시간을 받을 기회를 제공하지 않습니다. 이것이 바로 GPU 점유(GPU squatting)가 가능하게 만든 시스템 속성입니다. 또한 이는 온콜 엔지니어들이 지속적인 유지보수 문제를 해결하기 위해 장시간 실행되는 작업 소유자들과 협상하도록 강요했던 원인이기도 합니다.

이러한 문제들을 해결하기 위해, 우리는 '스케줄링 계약(scheduling contract)'을 도입했습니다. 클러스터 접근 권한과 교환하여, 워크로드는 최소 실행 시간 또는 의미 있는 진전을 이루는 데 필요한 가장 짧은 점유 시간을 선언해야 합니다. 이 기간 동안, 해당 워크로드는 선점(preemption)으로부터 보호받습니다. 이는 연구자에게 진전의 보장성을 제공하는 동시에, 스케줄러가 그 진전이 확보된 후 재조정할 권한을 가지며, 중단 가능한 워크로드를 자동으로 재대기열화(re-queueing)시킬 수 있게 합니다. 또는 사용자는 최소 실행 시간을 0으로 설정하여 GPU 시간이 할당되지 않아야 함을 나타낼 수도 있습니다. 이러한 워크로드는 항상 선점의 대상이 되지만, 어떤 예산에도 청구되지 않는다는 점에서 자유롭습니다.

워크로드 생명주기는 다음과 같은 패턴을 따릅니다:

• 최소 실행 시간과 재개 가능 여부를 지정하여 워크로드가 제출됩니다.
• 워크로드는 페어 셰어 알고리즘(fair-share algorithm)에 따라, 조회 기간(lookback window) 내 실제 점유율 대비 할당된 시간을 가중치로 받아 스케줄링됩니다.
• 워크로드는 최소 실행 시간 동안 실행되며, 이는 해당 할당량에 청구됩니다.
• 워크로드는 연관된 할당량이 다른 것들보다 우선순위를 부여하는 한 계속 실행될 수 있습니다. 이 시간 역시 할당량에 청구됩니다.
• 선점되어 재대기열화될 수 있으며, 이는 2단계로 돌아갑니다.
• 워크로드가 완료되면 모든 리소스에 대한 권리를 해제합니다.

이러한 약속들이 결합하여 우리의 스케줄러에 시간 분할(time-slicing)을 추가합니다. 실행 중인 워크로드는 자동으로 제거되고 재대기열화될 수 있어, 페어 셰어가 수렴하도록 하고 점유(squatting)를 방지합니다. 또한, 건강하지 않은 호스트들이 최소 실행 시간에 도달함에 따라 워크로드를 소진할 수 있게 하여, 복구 활동이 완전히 자동화되도록 합니다. 이 마지막 지점은 우리가 이 작업을 계획할 때 생각했던 것보다 더 중요했습니다. 이는 인간의 개입(human-in-the-loop)이 필요한 수리 작업량을 74% 감소시켰으며, 이는 온콜 업무 부담에서 엄청난 절감 효과를 가져왔습니다.

스케줄링 정책 변경은 의도치 않은 결과를 초래할 수 있다는 것을 알고 있습니다. 이 문제는 제로섬(zero-sum) 성격을 가지므로, 한 연구자에게 시간을 할당하는 것은 다른 연구자로부터 빼앗는 것을 의미합니다. 이런 교환을 잃은 사용자들은 새로운 우회 방법을 찾으려는 경향이 있습니다. 예산 기반 시스템을 도입하기 전에, 우리는 긴 대기 시간이 발생할 수 있는 지점을 빠르게 예측하고, 조회 기간(lookback window)의 길이 또는 최소 실행 시간을 허용하는 최대 값과 같은 설정 노브(configuration knobs)를 테스트하는 빠른 방법을 원했습니다 (우리는 8시간을 선택했습니다).