확장 가능한 LLM 학습을 위한 런타임 오케스트레이션 기반 2차 최적화 기술

2차 최적화 방법론 (Second-order methods)은 더 높은 샘플 효율성을 가진 LLM 학습을 향한 매력적인 경로를 제공하지만, 대규모 행렬 기반의 옵티마이저 상태 (optimizer states)를 유지하고 업데이트하는 데 드는 시스템 비용으로 인해 실제 사용에는 종종 제약이 따릅니다. 본 논문에서는 2차 최적화 로직을 핵심적인 GPU 학습 경로 (GPU training path)로부터 분리하여 이러한 병목 현상을 제거하도록 설계된 런타임 시스템인 extbf{Asteria}를 소개합니다. Asteria는 모든 프리컨디셔너 상태 (preconditioner state)를 가속기 (accelerator)에 유지하는 대신, 아키텍처 제약 조건과 런타임 압력에 따라 옵티마이저 상태를 GPU 메모리, CPU 메모리, 그리고 선택적인 NVMe 스토리지에 동적으로 분산합니다. 또한, Asteria는 트레이닝 훅 (training hooks)을 사용하여 섀도우 상태 (shadow states)를 사전에 준비하며, 이를 통해 GPU 연산이 계속되는 동안 호스트 (host)에서 비용이 많이 드는 역제곱근 (inverse-root) 계산을 비동기적으로 진행할 수 있게 합니다. 분산 학습 (distributed training)을 위해 Asteria는 토폴로지 인식 조정 (topology-aware coordination)을 통해 옵티마이저의 효과를 유지하면서 동기화 빈도를 제한하는 유계 지연 프로토콜 (bounded-staleness protocol)을 채택합니다. 우리는 메모리 제한 환경과 분산 학습 환경 모두에서 Asteria를 평가했습니다. 단일 GB10 GPU와 128GB 통합 메모리를 갖춘 DGX Spark 플랫폼에서 Asteria는 1B 파라미터 언어 모델에 대한 2차 학습을 지원합니다. 멀티 노드 GH200 시스템에서는 가시적인 옵티마이저 오버헤드를 낮추고, 반복적인 레이턴시 스파이크 (latency spikes)를 줄이며, 실제 시간 (wall-clock time) 기준 수렴을 가속화하고, 7B 파라미터 언어 모델에서 SOAP 및 KL-Shampoo의 최적화 이점을 유지합니다. 우리의 결과는 2차 LLM 학습이 단순히 옵티마이저를 단순화함으로써가 아니라, 옵티마이저 상태, 백그라운드 연산, 그리고 분산 동기화가 런타임 수준에서 어떻게 관리되는지를 재고함으로써 실용화될 수 있음을 시사합니다.