핵심 요약
대형 언어 모델과 멀티모달 모델의 복잡성이 증가함에 따라 기존 DeepSpeed의 제약적인 backward API와 높은 메모리 점유율이 개발의 병목이 되고 있다. Microsoft DeepSpeed 팀은 이를 해결하기 위해 PyTorch와 완전히 동일한 문법의 backward API를 도입하여 비스칼라(non-scalar) 그래디언트 처리를 지원하고 멀티모달 학습 루프를 간소화했다. 또한 파라미터와 그래디언트, 옵티마이저 상태를 모두 BF16/FP16으로 유지하는 저정밀도 학습 옵션을 추가하여 성능 저하 없이 최대 40%의 피크 메모리를 절감했다. 이러한 업데이트는 제한된 하드웨어 자원에서 더 큰 모델을 효율적으로 학습할 수 있는 환경을 제공한다.
배경
PyTorch, DeepSpeed, 분산 학습 기초, Mixed Precision 학습 개념
대상 독자
대규모 언어 모델(LLM) 및 멀티모달 모델을 학습하는 ML 엔지니어
의미 / 영향
이번 업데이트는 DeepSpeed를 순수 PyTorch 사용 경험에 더 가깝게 만들어 개발 생산성을 높이는 동시에 하드웨어 제약을 극복할 수 있는 메모리 효율성을 제공하여 더 넓은 범위의 연구자들이 고성능 모델 학습에 접근할 수 있게 한다.
섹션별 상세
PyTorch와 동일한 backward() 구문을 지원하여 복잡한 멀티모달 및 다중 컴포넌트 모델 학습 루프 작성이 훨씬 쉬워졌다. 기존 DeepSpeed 엔진은 스칼라 손실(loss)만 처리할 수 있었으나 이제는 PyTorch처럼 비스칼라 텐서에 대한 그래디언트 전달이 가능해져 비전 인코더와 LLM이 분리된 환경에서도 유연하게 학습할 수 있다.
분산 하이브리드 병렬 처리(Disaggregated Hybrid Parallelism)를 통해 멀티모달 모델 학습 속도를 30% 향상시켰다. Ray 액터 그룹을 사용하여 비전 인코더와 LLM을 별도로 관리하는 파이프라인에서 새로운 API를 적용한 결과 DeepSpeed-Ulysses와 같은 최적화 기술과 결합되어 높은 효율성을 입증했다. 이는 서로 다른 프로세스 간에 그래디언트를 주고받는 복잡한 구조에서도 DeepSpeed의 최적화 기능을 온전히 활용할 수 있음을 보여준다.
모델의 모든 상태인 파라미터, 그래디언트, 옵티마이저 상태를 BF16 또는 FP16 저정밀도로 유지하는 옵션이 추가되었다. 기존에는 수치적 안정성을 위해 FP32 마스터 파라미터를 유지해야 했으나 새로운 저정밀도 모드는 torch.autocast와 통합되어 수치적 안정성을 유지하면서도 7B 모델 기준 피크 메모리를 40% 절감하는 성과를 냈다.
Wikitext-103 데이터셋을 활용한 1000 스텝 학습 실험 결과 저정밀도 학습이 FP32 마스터 파라미터를 사용하는 베이스라인과 거의 동일한 손실(Loss) 수렴 곡선을 보여주었다. 최종 손실 값은 베이스라인 3.09 대비 저정밀도 3.12로 나타나 오차 범위 내의 성능을 유지했다. 이는 메모리 절감이 모델의 최종 품질이나 학습 안정성을 희생하지 않고도 달성될 수 있음을 시사하는 중요한 지표이다.
이미지 분석
1000 스텝 학습 동안 두 설정의 손실 값이 거의 동일하게 수렴함을 보여준다. 이는 저정밀도 학습이 메모리를 절감하면서도 모델 성능을 희생하지 않는다는 증거로 제시된다.
베이스라인과 BF16 저정밀도 학습의 손실 곡선 비교 그래프
</> 코드 예제 포함
실무 Takeaway
- 멀티모달 모델 개발 시 DeepSpeed의 새로운 backward API를 사용하면 PyTorch 순정 코드와 유사한 유연성을 확보하면서도 ZeRO 및 오프로딩 최적화를 그대로 누릴 수 있다.
- 메모리 부족 문제로 대형 모델 학습이 어려운 경우 DeepSpeed 설정에서 bf16_master_weights_and_grads를 활성화하여 피크 메모리를 최대 40%까지 줄일 수 있다.
- Ray와 DeepSpeed-Ulysses를 결합한 분산 하이브리드 병렬 처리를 적용하면 복잡한 모델 구조에서도 약 30%의 학습 속도 향상을 기대할 수 있다.
언급된 리소스
AI 분석 전체 내용 보기
AI 요약 · 북마크 · 개인 피드 설정 — 무료