Memory-Efficient Pipeline-Parallel DNN Training (PipeDream-2BW, PipeDream-Flush)

PipeDream의 메모리 부담을 개선

 

[Github]

[arXiv](Current version v3)

 

 

Abstract

메모리 효율적인 파이프라인 병렬성을 지원하는 시스템인 PipeDream-2BW 제안

 

 

Introduction

이전의 파이프라인 모델 병렬화:

[GPipe]

 

[PipeDream]

 

GPipe는 유휴 시간이 많으며, PipeDream은 가중치 저장으로 인해 메모리 제약이 크다는 단점이 있다.

 

PipeDream-2BW의 이점:

• 유휴 시간을 유발하는 pipeline flush를 피하면서 메모리 공간을 줄이는 double-buffered weight updates (2BW)
• 효과적인 병렬화 체계를 생성하는 PipeDream-2BW planner

 

 

 

PipeDream-2BW System Design

Double-Buffered Weight Updates (2BW)

 

기울기를 누적하여 m개의 마이크로 배치(이하 MB)마다 새로운 가중치 생성. 가중치 업데이트 시 업데이트 전 입력 MB는 여전히 이전의 가중치를 사용하므로 2개의 가중치 버전이 필요하다.

 

Memory Footprint

Worker 간의 경계에서만 활성화를 저장하고 나머지 활성화는 역전파에서 다시 계산하는 활성화 재계산을 통해 메모리 절약.

 

Semantics

일반적인 SGD(확률적 경사하강법)의 업데이트:

 

2BW의 업데이트:

예를 들어, MB8 역전파 후 업데이트를 하려고 할 때, 업데이트에 포함되는 MB 5,6,7(8까지 포함되는진 모르겠다)은  W⁽⁴⁾가 아닌 W⁽⁰⁾으로 전파되었기 때문이다.

 

그래도 모델의 피라미터가 단일 반복에서 별로 바뀌지 않기 때문에 성능에 차이는 없다고 한다.

 

Weight Updates with Flushes (PipeDream-Flush)

PipeDream with pipeline flush version. 2BW보다 처리량은 낮지만 가중치 저장이 없어 메모리 효율적이다.

 

Memory Footprint

PipeDream은 초기에 입력된 MB의 역전파가 빨리 끝나기 때문에 '활성'된 활성화의 수가 GPipe보다 평균적으로 적어 상대적으로 낮은 memory footprint를 제공한다. (최대 활성화 메모리는 같음)

2BW보다는 당연히 더 낮다.

 

Semantics

가중치 저장 없어서 일반 SGD와 같음.

 

Equi-replicated Stages (Parallel Pipelines)

파이프라인을 병렬할 수 있다.... 는 듯...?

 

이런 느낌...?

파이프라인 간의 통신 없이 모든 병렬 파이프라인에서 기울기를 집계하여 업데이트하면 된다고 한다.

 

 

 

Planner

가장 많은 처리량을 감당할 수 있는 stage 분할을 결정하는 알고리즘.

누군가가 메타 인지가 가장 중요하다고 했죠. 자세한 사항은 논문의 부록 A에 있습니다.