대규모 신경망 훈련을 위한 OpenAI 기술

대규모 신경망을 훈련하려면 단일 동기화 계산을 수행하기 위해 GPU 클러스터를 조정해야 합니다. 모델과 클러스터 규모가 커짐에 따라 실무자들은 계산 부하와 메모리 요구 사항을 여러 하드웨어 가속기에 분산시키기 위해 다양한 병렬 처리 기법을 사용합니다.

데이터 병렬 처리

데이터 병렬 처리는 동일한 모델 파라미터를 각 워커에 복사하고, 훈련 배치의 서로 다른 부분을 동시에 처리하도록 할당함으로써 여러 GPU를 사용할 수 있게 합니다. 이는 많은 GPU의 연산 능력을 활용할 수 있게 하지만, 모델은 여전히 단일 GPU 메모리 안에 들어가야 합니다.

워커 간 일관성을 유지하려면 시스템이 파라미터 업데이트를 조정해야 합니다. 표준 동기식 접근 방식은 세 단계로 이루어집니다:

각 워커가 독립적으로 그래디언트를 계산합니다.
모든 워커의 그래디언트를 평균화합니다(블로킹 통신 단계).
각 워커가 독립적으로 새로운 파라미터를 계산합니다.

블로킹 평균의 오버헤드를 줄이기 위한 비동기 동기화 방식도 존재하지만, 이는 학습 효율성을 떨어뜨리는 경우가 많아 대부분의 실무자는 동기식 방법을 선호합니다.

파이프라인 병렬 처리

파이프라인 병렬 처리는 모델의 연속적인 청크(연속 레이어)를 서로 다른 GPU에 나누어 배치함으로써 장치당 메모리 사용량을 줄입니다.

"버블" 문제 해결

단순 파이프라인 구현은 "버블"을 생성합니다—이전 장치의 출력이 도착하기를 기다리는 동안 워커가 유휴 상태에 머무는 기간입니다. 이러한 버블을 최소화하기 위해 배치를 더 작은 마이크로배치로 나눕니다. 각 워커는 마이크로배치가 사용 가능해지는 즉시 처리하기 시작해 계산과 대기 시간을 겹치게 합니다. 그래디언트는 마이크로배치별로 평균화되고, 파라미터 업데이트는 모든 마이크로배치가 완료된 후에만 수행됩니다.

스케줄링 전략

전방 및 역전파를 관리하기 위해 두 가지 주요 스케줄링 방식이 사용됩니다:

GPipe: 워커가 전방과 역전파를 순차적으로 처리하고, 여러 마이크로배치의 그래디언트를 마지막에 동기식으로 집계합니다.
PipeDream: 워커가 전방과 역전파를 교대로 처리합니다. 이는 효율성을 높일 수 있지만 일부 계산이 오래된 파라미터를 사용할 수 있습니다.

텐서 병렬 처리

텐서 병렬 처리는 레이어 내 개별 연산을 GPU들 사이에 "수평"으로 나눕니다. Transformer와 같은 최신 아키텍처에서는 활성화 배치 행렬과 큰 가중치 행렬을 곱하는 것이 주요 병목 현상입니다. 텐서 병렬 처리는 가중치 행렬을 균등한 조각으로 분할하고, 각 조각을 다른 GPU에 배치해 전체 행렬 곱의 일부를 계산한 뒤 결과를 통합합니다.

구체적인 구현 예시:

Megatron-LM: Transformer의 MLP와 self‑attention 레이어 내 행렬 곱을 병렬화합니다.
PTD-P: 텐서, 데이터, 파이프라인 병렬 처리를 결합하고, 연속되지 않은 레이어를 각 장치에 할당해 버블 오버헤드를 감소시킵니다.
Sequence Parallelism: 입력 시퀀스를 시간 축으로 여러 하위 예제로 나누어 피크 메모리 사용량을 줄입니다.

전문가 혼합 (Mixture‑of‑Experts, MoE)

Mixture‑of‑Experts (MoE)는 모델이 파라미터 수는 크게 늘리면서도 계산 비용은 비례적으로 증가하지 않게 합니다. 이는 입력마다 네트워크 가중치 중 일부("전문가")만 선택해 출력을 계산하는 게이팅 메커니즘을 사용함으로써 이루어집니다. 서로 다른 전문가를 서로 다른 GPU에 배치할 수 있기 때문에, MoE는 모델에 사용되는 GPU 수를 확장하는 스케일러블한 방법을 제공합니다.

메모리 절감 설계

병렬 처리 외에도 훈련에 필요한 장치 RAM을 줄이기 위한 여러 전략이 사용됩니다:

Checkpointing (Activation Recomputation): 그래디언트 계산을 위해 모든 활성화를 저장하는 대신, 체크포인트는 일부만 저장하고 역전파 중에 중간 활성화를 즉시 재계산합니다. 비용이 많이 드는 활성화만 체크포인트하고, 계산은 저렴한 활성화는 재계산하는 선택적 재계산이 이를 더욱 최적화합니다.
Mixed Precision Training: 낮은 정밀도(보통 FP16)로 훈련하면 FLOP 수는 증가하면서도 정확도 손실이 최소인 상태로 장치 RAM을 절약합니다.
Offloading: 사용되지 않는 데이터를 일시적으로 CPU나 다른 장치로 옮깁니다. ZeRO 구현은 파라미터, 그래디언트, 옵티마이저 상태를 하드웨어에 분산시키고 필요할 때만 메모리에 올립니다.
Memory Efficient Optimizers: Adafactor와 같은 옵티마이저를 사용해 실행 상태의 메모리 footprint를 감소시킵니다.
Compression: 중간 결과를 압축합니다. 예를 들어 Gist는 역전파를 위한 활성화를 압축하고, DALL·E는 동기화 전에 그래디언트를 압축합니다.