[yongggg's] FSDP: Fully Sharded Data Parallel

안녕하세요 이번 장에서는 LLM의 등장으로 모델 학습의 어려움이 커지면서 DDP의 병렬 분산 처리 학습으로 담을 수 없는 큰 모델을 학습하기 위한 방법인 Fully Sharedd Data Parallel(이하 FSDP)에 대해 설명을 드리겠습니다. 원문은 아래에 있으니, 아래를 참고하시는 것도 추천합니다!

https://engineering.fb.com/2021/07/15/open-source/fsdp

 

Large AI 모델을 쉽게 학습하기 위한 infra로는 intra-layer model parallelism, pipeline model parallelism, optimizer state+gradient sharding, and mixture of experts 등의 방법이 있습니다. 이 중에서 아래 그림에서도 볼 수 있듯이,  FSDP를 사용한다면 훨신 큰 모델을 보다 효율적으로 학습할 수 있습니다. 이는 FairScale library에서 구현이 되었으며, FSDP 기술은 이미 Facebook에서 이미 NLP 및 Vision models을 학습하는데 사용되고 있습니다. FSDP는 표준 DDP(Distributed Data Parallel) training 동일한 결과를 제공하며, PyTorch의 Distributed Data Parallel 모듈을 대체하는 사용하기 쉬운 인터페이스에서 사용할 수 있습니다. 그럼 자세한 설명을 시작해보겠습니다.

How FSDP Works

standard DDP training에서는 모든 workers에서 별도의 배치를 처리하고 all-reduce operation을 사용하여 모든 workers에 걸쳐 gradient를 합산한다. DDP는 매우 인기가 있지만 모델 weights와 optimizer states가 모든 DDP workers에 걸쳐 복제되기 때문에 필요한 것보다 더 많은 GPU memory가 필요하다.

 

복제를 줄이기 위한 한 가지 방법은 full parameter sharding이라는 process를 적용하는 것인데, local computation에 필요한 model parameters, gradients 그리고 optimizers들의 부분집합만을 사용할 수 있게 한다. 이 방법의 구현체인 ZeRO-3은 Microsoft에 의해 이미 대중화 되었다.

 

full parameter sharding을 unlock하기 위한 핵심 insight는 DDP의 all-reduce operations를 seperate reduce-scatter와 all-gather operations로 분해할 수 있다는 것이다. 이는 다음 그림과 같이 나타낼 수 있다.

all-reduce는 reduce-scatter와 all-gather의 조합이다. gradient를 집계하는 standard all-reduce 연산은 두 개의 별개의 단계, 즉 reduce-scatter와 all-gather로 분해될 수 있다. reduce-scatter 단계 동안 gradient들은 Rank Index를 기반으로 각 GPU에서 ranks 간 동일한 blocks으로 합산된다. all-gather 단계 동안 각 GPU에서 사용할 수 있는 aggregated(집계) gradients의 공유 부분은 모든 GPU에서 사용할 수 있게 된다.그런 다음 각 DDP workers가 parameters와 optimizer states의 단 하나의 shard만 저장하면 되도록 reduce-scatter와 all-gather를 재정렬할 수 있다.

 

아래 그림은 표준 DDP 훈련(위)과 FSDP 훈련(아래)을 보여준다:

standard data parallel training과 fully sharded data parallel training의 비교.

standard data parallel training methods에서 models의 복사본이 각 GPU 상에 존재하고 forward and backward 패스의 시퀀스가 데이터의 shard에 대해서만 평가된다. 이러한 Local computations 후에, global gradients update를 계산하기 위해, 각 Local process에 대한 parameter 및 optimizers가 다른 GPU들과 공유된다.

FSDP에서는 오직 모델의 shard만이 GPU 상에 존재한다. 그런 다음, loca에서 forward pass를 계산하기 위해 (all-gather step에 의해) 다른 GPU들로 부터 모든 가중치가 수집된다. 이러한 가중치의 수집은 이후 backward pass 전에 다시 수행된다. backward pass 이후에 Local gradients는 reduce-scatter 단계에 의해서 평균되고 GPU들에 걸쳐 sharding되며, 이는 각각의 GPU가 그것의 local weight shard를 업데이트할 수 있게 한다.

 

memory 효율을 극대화하기 위해 각 layer의 forward pass 후에 전체 가중치를 폐기하여 후속 layer의 memory를 절약할 수 있다. 이는 Network의 모든 Layer에 FSDP wrapper를 적용하여 구현할 수 있다.(reshard_after_forward=True).

 

In pseudo-code:

FSDP forward pass:
    for layer_i in layers:
        all-gather full weights for layer_i
        forward pass for layer_i
        discard full weights for layer_i

FSDP backward pass:
    for layer_i in layers:
        all-gather full weights for layer_i
        backward pass for layer_i
        discard full weights for layer_i
        reduce-scatter gradients for layer_i

How to use FSDP

FairScale의 FSDP 라이브러리는 대규모 훈련의 많은 중요한 측면들에 대한 낮은 수준의 옵션들을 노출하고 있다. FSDP를 전력으로 적용할 때 고려해야 할 몇 가지 중요한 영역들은 다음과 같다.

1. Model wrapping: 일시적인 GPU 메모리 요구를 최소화하기 위해, 사용자들은 모델을 중첩된 방식으로 랩핑할 필요가 있다. 이것은 추가적인 복잡성을 도입한다. auto_wrap 유틸리티는 중첩된 랩핑 목적으로 기존의 PyTorch 모델 코드에 주석을 달 때 유용하다.
2. Model initialization: DDP와 달리 FSDP는 GPU 작업자 간에 모델 가중치를 자동으로 동기화하지 않는다. 이는 모든 GPU 작업자가 동일한 초기 가중치를 갖도록 모델 초기화를 신중하게 수행해야 함을 의미한다.
3. Optimizer settings: 샤딩 및 래핑으로 인해 특정 유형의 Optimizer 및 Optimizer 설정만 FSDP에 의해 지원된다. 특히 모듈이 FSDP에 의해 래핑되고 그 파라미터가 단일 텐서로 평탄화되면 사용자는 이러한 모듈에서 상이한 파라미터 그룹에 대해 상이한 하이퍼 파라미터를 사용할 수 없다.
4. Mixed precision: FSDP는 FP16 마스터 웨이트를 사용하여 고급 혼합 정밀도 훈련을 지원할 뿐만 아니라 FP16을 그라디언트 상에서 reduce와 scatter 시킨다. 완전 정밀도를 사용하는 경우에만 모델의 특정 부분들이 수렴될 수 있다. 그러한 경우, 모델의 부분들을 완전 정밀도로 선택적으로 실행시키기 위해 추가적인 랩핑이 필요하다.
5. State checkpointing and inference: 모델 스케일이 크면 모델 상태를 저장하고 로드하는 것이 어려워질 수 있다. FSDP는 그 작업을 가능하게 하는 여러 가지 방법을 지원하지만 결코 사소하지 않다.
6. 마지막으로 FSDP는 FairScale의 checkpoint_wrapper와 같은 activation checkpointing 기능과 함께 사용되는 경우가 많다. 사용자는 제한된 GPU 메모리 공간 내에서 대규모 모델에 맞게 활성화 체크포인트 전략을 신중하게 조정해야 할 수 있다.