기술 요약: PrismLLM – 소수 GPU 로 충실한 LLM 학습 에뮬레이션 구현

1. 문제 제기

대규모 언어 모델 (LLM) 의 학습은 수천 개의 GPU 로 구성된 클러스터 규모로 확장되었습니다. 이러한 규모는 모델 발전을 주도하지만, 학습 프레임워크 자체의 개발, 디버깅, 최적화에 있어 상당한 병목 현상을 초래합니다. 엔지니어들은 다음과 같은 "DevOps 딜레마"에 직면해 있습니다:

• 자원 부족: 대부분의 GPU 가 프로덕션 워크로드에 할당되어 실험을 위한 자원이 부족합니다. 디버깅을 위해 수천 개의 GPU 를 전용으로 할당하는 것은 비용이 너무 많이 들고 운영상 방해가 됩니다.
• 기존 솔루션의 실패:
  • 시뮬레이션: 전통적인 시뮬레이터는 복잡한 분석 모델과 실행 추적을 의존합니다. 학습 스택이 프레임워크, 컴파일러, 하드웨어 세대에 따라 빠르게 진화함에 따라 이러한 모델의 충실도를 유지하는 것이 점점 더 어렵고 비용이 많이 듭니다.
  • 축소 실행: 전체 소프트웨어 스택을 소수의 GPU 에서 실행하는 것은 규모에 의존하는 동작을 포착하지 못합니다. 월드 사이즈, 병렬화 전략 (TP, PP, EP), 그리고 배치 크기의 변화는 통신 구조, 메모리 레이아웃, 병목 현상을 근본적으로 변경합니다. 결과적으로 축소 실험에서 도출된 결론은 종종 프로덕션 환경으로 이전되지 않습니다.

핵심적인 과제는 전체 규모 배포의 비용을 치르지 않고 대규모 학습 동작을 충실하게 연구하는 것입니다.

2. 방법론: PrismLLM

PrismLLM 은 논리적 학습 규모와 물리적 GPU 할당을 분리하도록 설계된 시스템입니다. 이를 통해 엔지니어들은 소수의 물리적 GPU 만으로도 충실한 대규모 동작 하에서 특정 "관심 순위 (ranks of interest)"를 실행하고 관찰할 수 있습니다. 이 시스템은 두 가지 명확한 단계로 작동합니다:

1 단계: 고정밀 그래프 구축 (준비)

목표는 전체 규모의 클러스터 없이 모든 순위에 대한 연산, 통신, 종속성을 포착하는 글로벌 실행 그래프를 생성하는 것입니다.

• 컨텍스트 스위칭 실행: PrismLLM 은 논리적 순위를 소수의 물리적 GPU 에 멀티플렉싱하는 중앙 집중식 조정기를 사용합니다. 이 조정기는 통신 지점에서 차단될 때까지 순위 하위 집합을 실행하도록 스케줄링합니다.
• 상태 보존: 순위가 차단되면 실행 상태와 통신 컨텍스트가 CPU 메모리 (또는 디스크) 에 저장되고, GPU 는 다른 실행 가능한 순위로 교체됩니다. 이를 통해 시스템은 소수의 GPU 에서 대규모 작업 (예: 1024 개 이상의 순위) 의 전체 실행 그래프를 포착할 수 있습니다.
• PrismTrace: 실행 그래프를 기록하기 위해 새롭고 경량인 추적 형식인 PrismTrace 가 사용됩니다. PyTorch 프로파일러와 같은 표준 프로파일러가 세밀한 연산자 세부 사항을 포착하는 것과 달리, PrismTrace 는 스케줄링 단위 (예: 마이크로배치) 에 초점을 맞추고 필수적인 통신 타이밍과 종속성만 포착하여 오버헤드를 크게 줄입니다.
• 타이밍 보정: 초기 그래프 ("Bare Graph") 는 컨텍스트 스위칭 오버헤드로 인해 정확한 타이밍이 없습니다. PrismLLM 은 워크로드를 "슬라이스"로 나눕니다. 각 슬라이스에서 물리적 GPU 를 실제 연산에 사용하고 나머지는 가상 순위로 재생하여 슬라이스를 순차적으로 실행합니다. 그런 다음 종속성 정보를 사용하여 슬라이스 간 타이밍을 정렬하고 전역적으로 일관된 고정밀 실행 타임라인을 재구성하는 슬라이스 간 타이밍 보정을 수행합니다.

2 단계: 하이브리드 에뮬레이션

고정밀 그래프가 준비되면 PrismLLM 은 실제 에뮬레이션을 수행합니다.

• 하이브리드 실행: 선택된 "샌드박스" 순위가 물리적 GPU 에서 원래 학습 프로그램을 실행합니다. 나머지 순위는 "어시스턴트" GPU 에서 가상 순위로 인스턴스화됩니다.
• 가상 순위 재생: 가상 순위는 실제 연산을 수행하지 않습니다. 대신 사전 보정된 실행 그래프를 탐색하여 기록된 시간 동안 대기하고, 상호작용 구조를 보존하기 위해 샌드박스 순위와 실제 통신 작업을 수행합니다.
• 효율성을 위한 최적화:
  • 가상 순위 초기화: 수천 개의 NCCL 그룹을 초기화하는 막대한 오버헤드를 피하기 위해 PrismLLM 은 NCCL 그룹 축소를 사용하여 샌드박스 순위와 겹치는 그룹만 인스턴스화합니다.
  • 런타임 통신 가지치기: 자원 경쟁을 방지하기 위해 비인접 가상 순위는 집단 통신 경로에서 제거됩니다. PrismLLM 은 NCCL 의 데이터 전송 논리를 수정하여 (예: 링 올리듀스) 왼쪽 가장자리 가상 순위가 제거된 순위의 누락된 기여분을 보상하도록 함으로써 수치적 정확성을 보장합니다.
  • 텐서 프리페칭: 런타임 프리페칭 파이프라인은 집단 연산을 위한 텐서를 관리하여 유휴 기간 동안 GPU 버퍼 풀로 로드하여 경쟁을 방지합니다.

3. 주요 기여

1. 하드웨어로부터의 규모 분리: PrismLLM 은 논리적 규모를 가상화하는 패러다임을 도입하여 엔지니어들이 물리적 하드웨어의 일부만으로도 대규모 상호작용을 관찰할 수 있게 합니다.
2. 고정밀 그래프 생성: 이 시스템은 컨텍스트 스위칭과 슬라이스 간 보정을 통해 모든 순위 간의 종속성과 타이밍을 포착하는 완전한 실행 그래프를 구축하여 준비 단계에서 전체 규모 실행이 필요하지 않습니다.
3. 정확성 보장이 있는 하이브리드 에뮬레이션: 이 시스템은 일부 순위를 실제 하드웨어에서 실행하고 나머지는 가상 참여자로 재생합니다. 집단 연산이 수치적으로 정확하도록 하고 샌드박스 순위가 현실적인 통신 패턴을 경험하도록 보장하기 위해 NCCL 그룹 축소, 통신 가지치기, 수치적 보상과 같은 새로운 기법을 도입했습니다.
4. 제로 코드 수정: 이 시스템은 엔지니어들이 수정되지 않은 학습 코드를 실행하고 시뮬레이션이나 축소를 위해 모델을 다시 작성하는 엔지니어링 오버헤드 없이 기존 코드베이스를 재사용할 수 있게 합니다.

4. 실험 결과

PrismLLM 은 512 개에서 8,192 개에 이르는 클러스터 규모에서 대규모 LLM 학습 워크로드 (Qwen 3 MoE) 로 평가되었습니다.

• 정확도:
  • 반복 시간: 다양한 모델 크기와 병렬화 전략에서 엔드투엔드 반복 시간을 예측하는 데 평균 **0.58%**의 오차를 달성했습니다.
  • 메모리 사용량: 피크 GPU 메모리 사용량을 0.01% 미만의 오차로 재현했습니다. 베이스라인 실험에서 발생했던 메모리 부족 (OOM) 오류도 성공적으로 재현했습니다.
  • 커널 충실도: 에뮬레이션된 커널 지속 시간과 시작 시간은 베이스라인 실행에서 관찰된 자연스러운 하드웨어 변동 범위 내에 있었습니다.
• 효율성:
  • PrismLLM 은 원래 배포에 필요한 물리적 GPU 의 1% 미만 (구체적으로 8,192 개 GPU 타겟에 대해 32 개의 물리적 노드) 을 사용하여 8,192 개 GPU 클러스터를 에뮬레이션할 수 있습니다.
  • 어시스턴트 노드를 목표 파이프라인 병렬화에 비례하여 확장함으로써, 시스템은 목표 규모가 증가하더라도 선형 시간 오버헤드를 피하고 안정적인 에뮬레이션 시간 (80 분 미만) 을 유지합니다.
• 시뮬레이터와의 비교:
  • Phantora와 비교했을 때, PrismLLM 은 특히 다양한 시퀀스 길이와 복잡한 집단 커널에서 훨씬 더 높은 정확도를 보였으며, Phantora 의 오류는 최대 64% 까지 급증했습니다.
  • SimAI와 비교했을 때, PrismLLM 은 정확도 면에서 압도적으로 우수했습니다 (SimAI 는 약 77% 의 오류를 보임). 이는 주로 SimAI 가 파이프라인 병렬화에서의 순위 간 종속성과 MoE 고유의 오버헤드를 포착하지 못했기 때문입니다.

5. 의의 및 주장

이 논문은 PrismLLM 이 LLM 학습 라이프사이클의 중요한 격차, 즉 prohibitive 한 자원 비용 없이 충실한 대규모 시스템 실험을 수행할 수 없는 문제를 해결한다고 주장합니다.

• 실용적인 DevOps: 엔지니어들이 작은 클러스터에서 최적화를 빠르게 반복하고, 전체 부하에서만 나타나는 규모 의존적 실패 (예: 열적 스로틀링 문제) 를 디버깅하며, MoE 라우팅 균형과 같은 구성을 높은 확신으로 조정할 수 있는 워크플로우를 가능하게 합니다.
• 근사치가 아닌 충실도: 동작을 근사화하는 시뮬레이션과 달리, PrismLLM 은 관심 순위에서 실제 하드웨어에서 실제 프로그램을 실행하여 관찰된 동작 (타이밍, 메모리 압력, 동기화) 이 현실에 기반하도록 보장합니다.
• 확장성: 이 시스템은 최소한의 하드웨어 발자국으로 수천 개의 GPU 를 가진 대규모 클러스터를 에뮬레이션할 수 있음을 보여주어, 프로덕션 자원이 완전히 할당된 상황에서도 대규모 시스템 디버깅을 접근 가능하게 만듭니다.

저자들은 PrismLLM 과 관련 추적 데이터를 오픈소스하여 추가 연구 및 산업 채택을 촉진하겠다고 약속합니다.