DeepSpeed 업데이트: PyTorch 호환 Backward API 및 저정밀도 모델 학습 지원

섹션별 상세

PyTorch와 동일한 backward() 구문을 지원하여 기존 DeepSpeed 엔진 API의 제약을 극복했다. 이전에는 스칼라 손실값만 처리 가능했으나, 이제는 PyTorch 네이티브 방식처럼 여러 모델과 손실 함수를 조합하거나 비스칼라 텐서에 대한 그래디언트 전파가 가능하다.

멀티모달 모델 학습을 위한 분산 하이브리드 병렬 처리(Disaggregated Hybrid Parallelism)가 가능해졌다. Ray 액터 그룹을 사용하여 비전 인코더와 LLM을 별도로 처리할 때, LLM에서 비전 인코더로 비스칼라 그래디언트를 전달하는 복잡한 루프를 PyTorch와 유사한 깔끔한 코드로 구현할 수 있다.

python

# Runs on LLM actors
def text_backward_step(self):
    self.loss.backward()
    return self.vision_embeddings.grad.detach().clone()

# Runs on Vision actors
def vision_backward_step(self, vision_embedding_grad):
    self.vision_output.backward(gradient=vision_embedding_grad)

멀티모달 모델에서 비스칼라 그래디언트를 전달하는 PyTorch 호환 backward API 사용 예시

모든 모델 상태(파라미터, 그래디언트, 옵티마이저 상태)를 BF16 또는 FP16 저정밀도로 유지하는 기능을 도입했다. 기존에는 FP32 마스터 가중치를 유지해야 했으나, 이제는 이를 생략하고 torch.autocast와 결합하여 메모리 효율을 극대화하면서도 학습 수렴 성능을 유지한다.

json

{
  "zero_optimization": {
    "stage": 3
  },
  "bf16": {
    "enabled": true,
    "bf16_master_weights_and_grads": true,
    "bf16_optimizer_states": true
  },
  "torch_autocast": {
    "enabled": true,
    "dtype": "bfloat16"
  }
}

모든 모델 상태를 BF16 저정밀도로 유지하도록 설정하는 DeepSpeed 구성 예시

7B 모델과 ZeRO-3를 사용한 실험 결과, 저정밀도 학습 적용 시 피크 메모리가 31.38GB에서 18.93GB로 약 40% 감소했다. Wikitext-103 데이터셋을 통한 1000 스텝 학습 테스트에서 FP32 마스터 가중치 방식과 거의 대등한 최종 손실값(3.09 vs 3.12)을 기록하며 수치적 안정성을 입증했다.

Baseline(FP32 마스터)과 BF16 저정밀도 학습 간의 손실(Loss) 곡선 비교 그래프 — Chart1000 스텝 동안 두 방식의 손실 값이 거의 일치하며 겹쳐서 진행되는 것을 보여준다. 이는 저정밀도 학습이 메모리를 대폭 절감하면서도 모델의 수렴 성능과 수치적 안정성을 FP32 방식 수준으로 유지함을 증명한다.

용어 해설

역전파(Backward Pass): — 신경망 학습 과정에서 출력층의 오차를 입력층 방향으로 거꾸로 전파하며 각 가중치에 대한 그래디언트를 계산하는 단계이다. 모델의 파라미터를 업데이트하기 위한 필수 과정으로, 복잡한 모델 구조에서는 이 흐름을 제어하는 유연성이 중요하다.
Bfloat16(BF16): — FP32와 동일한 지수 범위를 가지면서 크기는 16비트로 줄인 부동소수점 형식이다. 딥러닝 학습 시 FP16보다 수치적 안정성이 뛰어나며, 메모리 사용량을 절반으로 줄이면서도 모델 성능 저하를 최소화할 수 있어 널리 사용된다.
ZeRO 최적화(ZeRO): — 모델 파라미터, 그래디언트, 옵티마이저 상태를 여러 GPU에 분산하여 저장함으로써 메모리 중복을 제거하는 기술이다. 대규모 언어 모델 학습 시 단일 GPU 메모리 한계를 극복하고 더 큰 모델을 학습할 수 있게 해주는 핵심 기법이다.
자동 혼합 정밀도(torch.autocast): — PyTorch에서 제공하는 기능으로, 연산의 특성에 따라 FP32와 저정밀도(FP16/BF16) 타입을 자동으로 선택하여 실행한다. 연산 속도를 높이고 메모리를 절약하면서도 중요한 계산의 정밀도를 유지하여 학습 안정성을 보장한다.
비스칼라 텐서(Non-scalar Tensor): — 단일 숫자값이 아닌 벡터, 행렬 또는 고차원 배열 형태의 데이터를 의미한다. 멀티모달 학습에서 한 모델의 출력이 다른 모델의 입력으로 쓰일 때, 역전파 과정에서 단순 손실값이 아닌 다차원 그래디언트를 전달해야 하는 경우에 사용된다.

기술

DeepSpeed
PyTorch
Ray
BF16
ZeRO-3

활용 사례

멀티모달 모델(비전 인코더 + LLM) 학습
제한된 GPU 메모리에서의 대형 모델 파인튜닝
분산 하이브리드 병렬 처리 시스템 구축

언급된 리소스

GitHubDeepSpeed GitHub Repository

문서DeepSpeed Configuration Documentation