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

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

핵심 요약

DeepSpeed는 복잡한 멀티모달 및 멀티 컴포넌트 모델 학습을 효율화하기 위해 두 가지 주요 업데이트를 발표했다. 첫째로 PyTorch와 동일한 backward() API를 지원하여 비스칼라 그래디언트 처리가 가능해졌으며, 이를 통해 멀티모달 학습에서 30%의 속도 향상을 달성했다. 둘째로 파라미터, 그래디언트, 옵티마이저 상태를 모두 BF16/FP16 저정밀도로 유지하는 옵션을 추가하여 수치적 안정성을 유지하면서도 피크 메모리를 40% 줄였다. 이러한 변화는 연구자들이 제한된 하드웨어에서 더 큰 모델을 더 유연하게 학습할 수 있게 돕는다.

배경

PyTorch 및 DeepSpeed 기본 사용법, ZeRO 최적화 개념, 혼합 정밀도(Mixed Precision) 학습에 대한 이해

대상 독자

대규모 언어 모델 및 멀티모달 모델을 학습시키는 ML 엔지니어 및 연구자

의미 / 영향

DeepSpeed의 이번 업데이트는 복잡한 모델 아키텍처를 PyTorch 네이티브 코드에 가깝게 구현할 수 있게 하여 개발 생산성을 높인다. 특히 저정밀도 마스터 가중치 지원은 하드웨어 제약이 있는 환경에서도 더 큰 모델을 학습할 수 있는 실질적인 해결책을 제시한다.

섹션별 상세

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 방식 수준으로 유지함을 증명한다.

실무 Takeaway

멀티모달 모델처럼 복잡한 그래디언트 흐름이 필요한 경우, 새로운 PyTorch 호환 API를 사용해 코드 복잡도를 낮추고 학습 속도를 최대 30% 향상시킬 수 있다.
메모리 부족으로 대형 모델 학습이 어려운 환경에서는 bf16_master_weights_and_grads 옵션을 활성화하여 성능 저하 없이 메모리 점유율을 40% 절감할 수 있다.
저정밀도 모델 상태 옵션을 사용할 때는 torch.autocast를 함께 활성화해야 수치적 안정성을 확보하고 학습 수렴을 보장할 수 있다.

언급된 리소스

GitHubDeepSpeed GitHub Repository

문서DeepSpeed Configuration Documentation