기술 요약: Kareus – 대규모 모델 학습에서의 동적 및 정적 에너지의 결합된 감소

문제 정의

AI 컴퓨팅 수요의 급격한 성장은 에너지 공급 속도를 앞질렀으며, 이로 인해 에너지는 대규모 모델 학습에 있어 비용이 많이 들고 경쟁이 치열한 자원이 되었습니다. 현재의 최적화 노력은 일반적으로 동적 에너지(활성 연산에 의해 소비되며, 주파수 스케일링을 통해 줄일 수 있음) 또는 정적 에너지(하드웨어 활동 여부와 관계없이 소비되며, 실행 시간을 단축함으로써 줄일 수 있음) 중 하나만을 다룹니다.

기존 솔루션들은 이러한 요소들을 독립적으로 최적화하거나 단순하게 결합한다는 결정적인 한계를 가지고 있습니다.

• 주파수 스케일링 접근 방식(예: Perseus)은 비임계 경로(non-critical paths)의 GPU 주파수를 낮추어 동적 에너지를 줄이지만, 순차적인 커널 실행을 따르는 경우가 많아 통신과 연산을 중첩(overlap)할 기회를 놓칩니다.
• 세밀한 커널 스케줄링 접근 방식(예: Nanobatching)은 통신과 연산을 중첩하여 정적 에너지를 줄이지만, 동적 에너지와 주파수 스케일링은 무시합니다.

본 논문은 SM 할당, 커널 런칭 타이밍, 그리고 GPU 주파수가 상호 의존적임을 지적합니다. 한 가지 요소를 변경하면 다른 요소들의 최적 구성이 바뀝니다. 예를 들어, 주파수를 낮추면 커널이 상대적으로 연산 집약적(compute-bound)으로 변하며, 이는 최적의 자원 경합 및 중첩 전략을 변화시킵니다. 결과적으로, 이러한 요소들의 일부만을 최적화하거나 단순하게 결합하는 것은 시간-에너지 트레이드오프 프런티어를 이론적 한계치까지 밀어붙이는 데 실패합니다.

방법론

저자들은 SM 할당, 커널 런칭 타이밍, GPU 주파수를 공동으로 최격화하여 파레토 최적(Pareto-optimal)의 시간-에너지 프런티어를 찾는 학습 시스템인 Kareus를 제안합니다.

1. 분할 중첩 실행 모델 (Partitioned Overlap Execution Model)

전역적인 문제를 해결 가능한 수준으로 만들기 위해, Kareus는 분할 중첩 실행 모델을 사용하여 전역 문제를 로컬 하위 문제들로 분해합니다.

• 분할(Partitioning): 순전파(forward) 및 역전파(backward)의 실행 그래프를 고정된 파티션으로 나눕니다. 하나의 파티션은 하나의 통신 커널(예: AllReduce)과, 해당 통신 커널과 데이터 의존성이 없는 다른 마이크로배치로부터 온 가장 긴 연속적인 연산 커널 시퀀스로 구성됩니다.
• 일반화(Generalization): 이 모델은 나노배칭(nanobatching) 기술을 일반화한 것으로, 모든 파티션에 대해 SM 할당 및 런칭 타이밍에 대한 세밀한 제어를 추가합니다.
• 계층적 구성(Hierarchical Composition): 시스템은 각 파티션을 독립적으로 최적화한 후, 이러한 로컬 프런티어들을 계층적으로 결합하여 글로벌 마이크로배치 및 이터레이션 레벨의 프런티어를 생성합니다.

2. 다중 패스 다중 목적 베이지안 최적화 (Multi-Pass Multi-Objective Bayesian Optimization, MBO)

설정 공간(SM 개수, 런칭 타이밍, 주파수 포함)에 대한 전수 조사는 비용이 매우 많이 들고(수천 GPU 시간이 소요되며 열 간섭을 유발함) 불가능하기 때문에, Kareus는 특화된 MBO 알고리즘을 채택합니다.

• 대리 모델(Surrogate Models): 실행 시간과 동적 에너지를 각각 예측하기 위해 두 개의 별도 그래디언트 부스팅 결정 트리(XGBoost) 모델을 사용합니다.
• 다중 패스 후보 선택(Multi-Pass Candidate Selection): 파레토 프런티어를 효과적으로 탐색하기 위해 여러 획득 함수(acquisition functions)를 병렬로 사용합니다.
  • 하이퍼볼륨 개선(Hypervolume Improvement, HVI): 세 가지 서로 다른 HVI 함수가 총 에너지 감소(낮은 주파수 선호), 낮은 동적 에너지(낮은 주파수 선호), 낮은 정적 에너지(빠른 실행 선호)를 향한 탐색을 가이드합니다.
  • 불확실성 탐색(Uncertainty Exploration): 대리 모델의 붓스트랩 앙상블을 통해 예측 불확실성을 추정하고, 탐색되지 않은 영역의 후보를 우선시합니다.
• 열적으로 안정적인 프로파일링(Thermally Stable Profiling): 정확한 에너지 측정을 보장하기 위해, 시스템은 5초간의 윈도우 동안 파티션을 실행하고 설정 간에 5초간의 냉각 기간을 삽입하여 열적 편향(thermal bias)을 방지합니다.

3. 시스템 구현

Kareus는 Megatron-LM 및 Perseus와 통합되어 구현되었습니다.

• 세밀한 SM 할당 제어를 위해 **MSCCL++**를 사용합니다.
• 파티션을 래핑하기 위해 커스텀 torch.autograd.Function을 사용하여 순전파와 역전파에 대해 서로 다른 스케줄을 허용합니다.
• 런타임 시, 사용자의 제약 조건(예: 시간 데드라인 또는 에너지 예산)에 따라 사전 계산된 프런티어에서 실행 스케줄을 선택하고, Perseus에서 변형된 컨트롤러를 통해 GPU 주파수를 동적으로 조정합니다.

주요 결과

Kareus는 Llama 3.2 3B 및 Qwen 3 1.7B에 대한 실제 학습을 포함한 14개의 워크로드와 16개의 NVIDIA A100 GPU를 이용한 Llama 3.3 70B의 대규모 에뮬레이션을 통해 평가되었습니다.

• 시간-에너지 트레이드오프: 최첨단 베이스라인(Megatron-LM + Perseus)과 비교하여 Kareus는 다음을 달성했습니다.
  • 동일한 학습 시간에서 최대 28.3%의 학습 에너지 감소.
  • 동일한 에너지 소비량에서 최대 27.5%의 학습 시간 단축.
• 최대 처리량 성능: 최대 처리량 모드에서, Kareks는 베이스라인인 Megatron-LM 대비 이터레이션 시간을 최대 14.9% 단축하고 에너지를 최대 22.1% 절감했습니다.
• 절제 연구(Ablation Studies): 주파수 스케일링이나 커널 스케줄링 중 하나라도 제거하면 상당한 성능 저하가 발생합니다(예: 커널 스케줄링을 제거하면 에너지가 10.8% 증가함). 이는 결합 최적화의 필요성을 확인시켜 줍니다.
• 확장성: Llama 3.3 70B에 대한 대규모 에뮬레이션 결과는 물리적 테스트베드 결과와 일관된 경향을 보여, 대규모 모델에 대한 시스템의 적용 가능성을 입증했습니다.

의의 및 주장

본 논문은 동적 에너지와 정적 에너지가 독립적으로 최적화될 수 없음을 입증함으로써, Kareus가 대규모 모델 학습을 위한 시간-에너지 트레이드오프 프런티어를 밀어올린다고 주장합니다. 핵심 기여는 실행 스케줄(SM 할당 및 타이밍)과 GPU 주파수가 결합되어 있으며 상호 의존적이라는 통찰입니다.

저자들은 다음을 강조합니다:

1. 결합 제어가 우월함: 주파수와 스케줄링을 결합하여 최적화하는 것이 주파수 제어를 하드웨어 전력 컨트롤러에 맡기거나 스케줄링을 주파수 인지 없이 최적화하는 것보다 입증 가능한 수준으로 더 에너지 효율적입니다.
2. 상호 의존성이 핵심임: 최적의 SM 할당과 런칭 타이밍은 주파수별로 다릅니다. 특정 주파수에서 최적인 스케줄이 다른 주파수에서는 최적이 아닐 수 있습니다.
3. 실용적 효율성: 문제를 파티션으로 분해하고 MBO를 사용함으로써, Kareus는 과도한 탐색 비용 없이 이러한 이득을 얻을 수 있어 현재 및 미래의 대규모 AI 학습 워크로드에 실용적인 솔루션이 됩니다.