🤖 machine learning

Runtime-Orchestrated Second-Order Optimization for Scalable LLM Training

이 논문은 최적화기 상태를 이종 메모리 계층 구조에 걸쳐 동적으로 분배하고 비싼 전구조건부 계산들을 중요한 GPU 훈련 경로에서 분리하여 오버헤드를 줄이고 수렴을 가속화함으로써 대규모 언어 모델을 위한 실용적이고 확장 가능한 2 차 최적화를 가능하게 하는 런타임 시스템인 Asteria 를 소개합니다.

원저자: Yishun Lu, Junhao Zhang, Zeyu Yang, Wes Armour

게시일 2026-05-18

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Yishun Lu, Junhao Zhang, Zeyu Yang, Wes Armour

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

거대하고 초지능적인 로봇(대형 언어 모델)에게 이야기 쓰기를 가르치려 한다고 상상해 보세요. 이를 위해 로봇의 실수를 보고 개선 방법을 알려주는 '교사'(최적화기) 가 필요합니다.

현재 대부분의 교사는 AdamW라는 간단한 방법을 사용합니다. 이는 숙제를 확인하고 몇 가지 틀린 답을 발견한 뒤 작고 빠른 수정을 가하는 학생과 같습니다. 빠르고 효과적이지만, 정말 잘하기 위해서는 (수백만 개의 예시와 같은) 엄청난 연습이 필요합니다.

2 차 최적화 (Second-Order Optimization) (SOAP 또는 Shampoo 와 같은) 라는 더 지능적이고 고급스러운 교사 방법이 있습니다. 이 교사는 단순히 무엇이 잘못되었는지 보는 것을 넘어, 실수의 형태를 분석하여 문제를 깊이 있게 이해합니다. 훨씬 더 빠르게 학습하며 더 적은 예시만으로도 충분합니다. 하지만 엄청난 단점이 있습니다. 이 똑똑한 교사는 압도적으로 무겁습니다. 복잡한 수학을 수행하기 위해 막대한 양의 메모리와 연산 능력이 필요하여, 종종 컴퓨터가 충돌하거나 너무 느려져 단순한 교사가 먼저 작업을 끝내기도 합니다.

Asteria 등장: '지능형 물류 관리자'

옥스퍼드 연구팀이 Asteria라는 새로운 시스템을 개발했습니다. Asteria 는 새로운 교사가 아니라, 똑똑한 교사가 학교를 망치지 않고 제 역할을 할 수 있도록 교실 전체를 재편성하는 천재적인 물류 관리자로 생각하세요.

Asteria 가 세 가지 가장 큰 문제를 어떻게 해결하는지 간단한 비유로 설명해 드리겠습니다.

1. "작은 방에 있는 가구" 문제 (메모리)

문제: 똑똑한 교사는 거대하고 복잡한 차트 (행렬) 를 머릿속 (GPU 메모리) 에 보관해야 합니다. 표준 컴퓨터에는 공간이 부족합니다. 스튜디오 아파트에 킹사이즈 침대, 식탁, 옷장을 모두 넣으려 하는 것과 같습니다. 방이 꽉 차면 교사는 일을 그만두어야 합니다.
Asteria 의 해결책: Asteria 는 접이식 가구 전문가처럼 행동합니다. 교사가 한 번에 모든 것을 손에 들고 있을 필요가 없다는 것을 깨닫습니다.

가장 활발하게 사용되는 차트 부분은 GPU (책상) 에 보관합니다.
무겁고 덜 자주 사용되는 부분은 CPU (복도) 나 심지어 하드 드라이브 (차고) 로 이동시킵니다.
결정적으로, 무거운 부분들은 정확히 필요할 때만 책상으로 다시 가져옵니다. 이를 통해 똑똑한 교사는 슈퍼컴퓨터뿐만 아니라 작은 노트북에서도 똑같이 잘 일할 수 있습니다.

2. "교통 체증" 문제 (속도)

문제: 몇 단계마다 똑똑한 교사는 차트를 업데이트하기 위해 거대하고 복잡한 계산 (거대한 퍼즐을 푸는 것과 같은) 을 수행해야 합니다. 이는 시간이 오래 걸려 전체 클래스의 작업을 멈추게 합니다. 단일 택배 상자를 내리기 위해 고속도로 전체를 막는 배송 트럭과 같습니다. GPU (컴퓨터의 두뇌) 는 계산을 기다리며 유휴 상태로 앉아 있습니다.
Asteria 의 해결책: Asteria 는 병렬 조립 라인을 도입합니다.

무거운 계산을 위해 메인 클래스를 멈추는 대신, Asteria 는 '무거운 작업'을 뒷사무실 (CPU) 에 있는 작업 팀에게 보냅니다.
메인 클래스 (GPU) 가 로봇을 가르치는 동안, 뒷사무실 작업자들은 조용히 배경에서 복잡한 퍼즐을 풉니다.
클래스가 새로운 차트가 필요할 때쯤이면, 뒷사무실은 이미 작업을 마치고 차트를 밀어 넣은 상태입니다. 메인 클래스는 멈출 필요가 없습니다. '교통 체증'이 사라집니다.

3. "그룹 채팅" 문제 (분산 학습)

문제: 여러 컴퓨터로 동시에 학습할 때, 모든 사람이 다음 단계로 넘어가기 전에 반드시 동일한 정보를 정확히 합의할 때까지 멈춰야 합니다. 이는 가장 느린 사람이 답장할 때까지 기다려야 모든 사람이 다시 타이핑할 수 있는 그룹 채팅과 같습니다.
Asteria 의 해결책: Asteria 는 "충분히 좋은" 정책을 사용합니다.

모든 사람이 완벽하게 동기화될 때까지 기다리는 대신, 컴퓨터들이 잠시 동안 약간 '오래된 (stale)' 정보를 가지고 작업할 수 있도록 허용합니다.
정말 필요할 때만 업데이트를 전송합니다.
이는 전체 프로젝트를 5 분마다 멈추는 대신, 서로를 신뢰하며 작업을 계속하다가 나중에 업데이트를 받아보는 팀처럼 그룹이 빠르게 움직이게 합니다.

결과: 그들은 무엇을 발견했나?

연구팀은 Nvidia DGX Spark 와 같은 강력한 단일 컴퓨터와 Nvidia GH200 과 같은 대규모 컴퓨터 클러스터를 포함한 실제 하드웨어에서 Asteria 를 테스트했습니다.

작은 기계에서도 작동합니다: 이전에는 이 똑똑한 교사의 메모리 요구 사항을 처리하지 못했던 단일 기계에서도 10 억 개의 파라미터를 가진 대형 언어 모델을 학습시킬 수 있었습니다.
실시간으로 더 빠릅니다: '교통 체증'을 숨기기 때문에, 수학이 더 어렵더라도 실제 시계 시간은 더 단축되어 학습이 완료됩니다.
에너지를 절약합니다: 컴퓨터의 두뇌 (GPU) 를 유휴 상태로 두지 않고 바쁘게 유지함으로써, Asteria 는 이전의 둔한 방법들보다 동일한 결과를 얻는 데 실제로 더 적은 총 에너지를 사용합니다.
품질이 떨어지지 않습니다: 모델은 '네이티브'(최적화되지 않은) 똑똑한 교사와 똑같이 잘 학습하지만, 훨씬 더 빠르고 저렴하게 도달합니다.

요약:
Asteria 는 새로운 수학 공식을 발명하지 않습니다. 대신 컴퓨터가 그 수학을 어떻게 실행하는지 재고합니다. 무거운 작업을 메인 작업과 분리하고, 사용 가능한 모든 저장 공간을 현명하게 활용하며, 컴퓨터들이 비동기적으로 작업하도록 합니다. 이는 이론적으로는 훌륭하지만 실제로는 불가능한 방법을 차세대 AI 학습을 위한 실용적인 도구로 바꿉니다.

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

1. 문제 제기

2 차 최적화 방법 (예: Shampoo, SOAP, KL-Shampoo) 은 AdamW 와 같은 1 차 옵티마이저에 비해 대규모 언어 모델 (LLM) 학습에서 더 우수한 통계적 효율성과 더 빠른 수렴을 제공합니다. 그러나 시스템 수준의 심각한 장벽으로 인해 대규모 실용 배포는 방해받고 있습니다.

메모리 제약: 2 차 방법은 대규모 공분산 행렬 ( $L$ 및 $R$ ) 을 저장하고 조밀한 고유값 분해를 수행해야 하므로, $O(N^2)$ 의 메모리 사용량과 $O(N^3)$ 의 계산 비용이 발생합니다. 메모리 제약 환경 (예: 통합 메모리 아키텍처를 갖춘 단일 노드 워크스테이션) 에서 이러한 상태는 물리적 메모리를 소진하여 Out-of-Memory (OOM) 실패를 초래하거나 모델 크기와 배치 크기를 축소하도록 강요합니다.
지연 시간 급증: 역근 (inverse roots, 예: $L^{-1/4}$ ) 의 주기적 계산에는 3 차 비용의 행렬 연산이 포함됩니다. 기존 구현에서는 이러한 동기식 업데이트가 GPU 학습 파이프라인을 차단하여, 효율적인 분산 학습 (예: FSDP) 에 필수적인 세밀한 계산과 통신의 오버랩을 방해하는 심각한 주기적 지연 시간 급증을 유발합니다.
동기화 오버헤드: 현재 분산 2 차 방법은 모든 랭크에 걸친 경직된 동기식 글로벌 상태 업데이트 (AllGather) 에 의존합니다. 이는 빠른 노드 내부 통신을 느린 노드 간 동기화와 결합하여, 시스템이 비동기적이고 토폴로지를 인식하는 노후화 (staleness) 한계를 활용하는 것을 방지합니다.

DTensor 기반 샤딩과 같은 기존 솔루션은 리소스 제약이 있는 단일 노드 설정이나 글로벌 동기화의 통신 오버헤드가 이점을 상쇄할 때 이러한 문제를 해결하지 못합니다.

2. 방법론: Asteria 런타임

Asteria 는 PyTorch FSDP 위에 구축된 하드웨어 - 소프트웨어 공동 설계 런타임 시스템입니다. 세 가지 핵심 메커니즘을 통해 2 차 최적화 논리를 중요한 GPU 학습 경로와 분리합니다.

A. 아키텍처 적응형 비대칭 메모리 계층화

상태를 동질적인 벡터로 취급하는 1 차 옵티마이저와 달리, Asteria 는 2 차 상태의 비대칭 수명 주기를 인식합니다.

GPU: 학습 중 빈번히 업데이트되는 factor_matrices(크로네커 인자) 를 저장합니다.
호스트 (CPU/UVM): inv_factor_matrices(역 인자) 를 저장합니다. GPU 메모리 압력을 피하기 위해 CPU 접근 가능 메모리에 유지됩니다.
NVMe (선택 사항): 용량 제한이 있는 실행을 위한 스테이징 계층을 제공하여, 직렬화된 계산 또는 JIT 페이지 페이징을 통해 상태를 스왑 아웃할 수 있게 합니다.
메커니즘: 런타임은 호스트 측 팩터 행렬 스냅샷을 캡처하고, 비동기 CPU 워커를 배치하여 역근을 계산한 후 결과를 호스트 메모리에 다시 씁니다. 필요 시 GPU 는 Unified Virtual Memory (UVM) 뷰 또는 명시적 프리페칭을 통해 이에 접근하여 불필요한 호스트 - 장치 복사본을 방지합니다.

B. 후크 오케스트레이션 섀도우 상태 파이프라인

지연 시간 급증을 제거하기 위해 Asteria 는 $O(N^3)$ 역근 계산을 중요한 GPU 경로에서 제거합니다.

섀도우 스트림: 저우선순위 CUDA 스트림이 GPU 측 상태 스테이징 (프리페칭 및 페이지 페이징) 을 처리하여 주요 학습 스트림을 차단하지 않습니다.
비동기 실행: 고비용 역근 업데이트는 CPU 워커 스레드에서 비동기적으로 실행됩니다.
후크 통합: 경량 PyTorch 후크 (register_forward_hook, register_full_backward_pre_hook) 가 스케줄링 신호로 작용합니다. 이들은 준비된 상태 전송의 배수를 유발하고 곧 소비될 역 인자의 프리페칭을 트리거하여, GPU 가 메인 학습 루프를 정지시키지 않고 최신 프리컨디셔너를 갖도록 보장합니다.

C. 토폴로지 인식 유계 노후화 일관성

분산 학습의 경우, Asteria 는 경직된 글로벌 동기화를 선택적 일관성 프로토콜로 대체합니다.

노후화 추적: CoherenceRegistry 가 각 프리컨디셔너 블록의 신선도를 추적합니다. 구성된 노후화 예산을 초과하는 블록만 동기화됩니다.
계층적 동기화: 동기화는 두 단계로 발생합니다. 먼저 노드 로컬 프로세스 그룹 내에서, 그 다음 노드 대표 간에 발생합니다. 이는 노드 간 트래픽을 최소화하고 호스트 접근 가능 메모리에서 통신을 유지함으로써 불필요한 CPU-to-GPU 왕복을 방지합니다.
유계 배리어: 런타임은 CPU 작업이 진행 중일 때 프리컨디셔너 버전이 너무 노후화되면 옵티마이저가 완료될 때까지 기다리도록 보장하는 로컬 배리어를 강제하여, 신선도와 처리량을 균형 있게 맞춥니다.

3. 주요 기여

런타임 분리: Asteria 는 2 차 최적화를 알고리즘 단순화가 아닌, 런타임 수준에서 상태 관리, 백그라운드 계산 및 동기화를 재고함으로써 실용화할 수 있음을 입증합니다.
메모리 확장성: 기존 샤딩이 불가능한 단일 노드, 메모리 제약 하드웨어 (예: 128GB 통합 메모리를 갖춘 단일 GB10 GPU) 에서 2 차 학습을 가능하게 합니다.
지연 시간 숨김: 역근 계산을 호스트로 이동하고 섀도우 파이프라인을 사용하여 주기적 지연 시간 급증을 평탄화하여, 1 차 방법과 비교 가능한 부드러운 하드웨어 실행 프로파일을 달성합니다.
에너지 효율성: 시스템은 GPU 유휴 시간을 줄이고 호스트 측 계산의 오버랩을 최적화하여 단위 에너지당 더 많은 손실 감소를 달성함으로써 에너지 - 손실 트레이드오프를 개선합니다.

4. 실험 결과

저자들은 OLMo 학습 프레임워크를 사용하여 두 가지 플랫폼에서 OLMo 계열 모델 (660M, 1B, 7B 파라미터) 에 대해 Asteria 를 평가했습니다.

단일 노드 (DGX Spark, GB10 GPU):
- 처리량: Asteria-SOAP 및 Asteria-KL-Shampoo 는 네이티브 구현에서 관찰된 주기적 단계 시간 급증을 제거했습니다 (피크 단계 시간을 약 96 초에서 약 66.5 초로 감소).
- 에너지: 네이티브 2 차 방법은 AdamW 보다 총 SoC 에너지를 약 17-20% 더 소비했으나, Asteria 는 이 격차를 크게 줄였습니다 (예: Asteria-KL-Shampoo 는 AdamW 보다 약 7.9% 더 많은 에너지만 소비).
- 효율성: Asteria-KL-Shampoo 는 가장 높은 정규화된 손실 감소 효율 (3.3010) 을 달성하여 네이티브 2 차 방법과 AdamW 모두를 능가했습니다.
멀티 노드 (GH200 클러스터):
- 수렴: Asteria 변형은 네이티브 2 차 대응물과 가까운 수렴 궤적을 유지하며, 목표 손실 수준에 도달하는 월 클락 시간에서 AdamW 를 크게 능가했습니다.
- 노후화 허용: 노후화 예산 ( $S$ ) 을 1 에서 5 로 변경하는 실험은 최종 모델 품질에 미미한 영향을 주면서 학습 시간을 현저히 단축함을 보여주었습니다.
- 강한 확장성: 2 노드에서 16 노드로 확장되는 7B 모델에서 Asteria 는 네이티브 베이스라인보다 더 높은 속도 향상과 더 낮은 단계당 시간을 달성하여, 상태 유지의 상대적 비용이 증가함에도 불구하고 효과적임을 입증했습니다.

5. 중요성 및 주장

이 논문은 실용적인 2 차 LLM 학습의 주요 장애물이 곡률 인식 업데이트의 수학적 비용이 아니라, 2 차 상태 유지와 현대 학습 스택 간의 시스템 불일치라고 주장합니다.

문제 재정의: 저자들은 2 차 최적화를 알고리즘, 런타임 스케줄링 및 이종 메모리 시스템을 포함하는 공동 설계 문제로 간주해야 한다고 주장합니다.
실현 가능성: 그들은 대규모 집계 GPU 메모리가 2 차 학습의 근본적인 전제 조건이 아님을 입증합니다. 수명 주기 인식 배치와 단계적 접근을 통해 이러한 방법은 엄격한 메모리 예산에서도 실행 가능합니다.
실용성: 프리컨디셔너 신선도와 통신을 경직된 글로벌 상태가 아닌 런타임 관리 자원으로 취급함으로써, Asteria 는 2 차 최적화의 통계적 약속과 실제 배포 가능성 간의 격차를 좁힙니다.

이 논문은 Asteria 가 알고리즘 단순화 없이 더 나은 수렴과 에너지 효율성을 제공함으로써 2 차 최적화를 LLM 학습을 위한 실현 가능한 시스템 선택지로 만든다고 결론지었습니다.