FSDP는 다음과 같이 동작한다.
- FSDP Unit - Vertically Splitting
- Sharding - Horiontally Splitting
- All Gather
- Reduce-Scatter
1. FSDP Unit - Verticall Splitting
FSDP Unit은 모델을 쪼개는 단위로 예를 들면 Layer, Stage, 또는 Layer의 Group(nn.Module)이 될 수 있다. 예를 들어 다음과 같이 34개의 Unit으로 구성 가능하다.
이 그림에서 FSDP Unit0는 Layer0, Layer3로 Unit1은 Layer1, Layer2로, Unit2는 Layer4, Layer5로 구성된다.
2. Sharding - Horizontal Splitting
FSDP Unit을 FlatParameter로 저장하고 FlatParameter를 각각의 GPU에 분할한다. (Sharding)
예를 들어 FSDP Unit이 12개의 Weight와 3개의 Bias를 가지고 있다고 하면 16개의 GPU로 분산한다고 하면 Padding을 하나 붙여서 GPU로 분할(Sharding)한다. Sharding 전략은 다음과 같이 사용 가능하다(DeepSpeed 그림)
이 그림에서 Baseline은 DDP와 같이 하나의 GPU에 모든 모델의 가중치가 있는 것이고, Pos는 Optimizer State를, Pos+g는 Optimizer State와 Gradient를, Pos+g+p는 이에 Parameter까지 분할하는 것이다. FSDP는 이중 FULL_SHARD (Parameter+Gradients+Optimizer States)와 SHARD_GRAD_OP(Gradients + Optimizer States), NO_SHARD(DPP처럼 동작), HYBRID_SHARD (하나의 노드에서는 FULL_SHARD, 노드간에는 Replcate) 하는 방식을 지원한다.
3. All-Gather (Both Forward + Backward)
NCCL이 지원하는 기능으로, 다음 그림과 같이 각 GPU에 분산되어 있는 데이터를 각각의 GPU가 전체 데이터를 합쳐주는 것이다.
FSDP에서는 다음과 같이 동작한다.
4개의 GPU에 Weight가 각각 Sharding되어 있는데, Forward와 Backward를 위해서 Unit에 있는 Weight를 각각의 GPU에 모아서 계산을 하고 계산 후에는 원래 Shard되어 있는 Weight만 남기고 해제하는 방식이다. (여기에서 중요한 점은 FSDP-Unit별로 All-Gather한다는 것이다.) 여기에서 배치는 각 GPU에 들어가는 배치에 GPU개수 만큼이 계산된다.
따라서, FSDP의 메모리 요구는 각 GPU에 저장되어야 하는 Sharded Model의 크기 + FSDP Unit이 합쳐졌을 때의 메모리 크기에 비례한다는 것이다.
상기1.의 예에서 FSDP는 다음과 같이 중첩해서 커뮤니케이션과 연산을 진행한다.
FSDP Unit0의 Layer0의 연산을 위해서 AG0로 파라미터를 각 GPU들로 모으고, 각각 Forward 진행(FWD0)하고 동시에 FSDP Uni1의 Layer1, Layer2 연산을 위해 AG1로 FSDP Unit1의 파라미터를 각 GPU로 모은다, 그리고 Layer1, Layer2의 Forward를 진행하고 동시에 FSDP Unit2의 파라미터를 각 GPU에 합친다. FSDP Unit1의 Forward연산이 끝나면 FSDP Unit1의 합쳐진 Parameter를 Free 시키고, FSDP Unit0의 Layer3에 Forward를 진행한다. 이 그림에서 보면 AG2다음에 또 AG2가 있는데 이는 Backward를 위한 것으로 4. Reduce-Scatter에서 설명한다.
4. Reduce-Scatter
Reduce-Scatter또한 NCCL의 기능중 하나로, 다음의 그림과 같이 각각의 GPU에 있는 데이터를 합치고 나눠서 분산 시키는 기능이다.
FSDP에서는 이 Feature를 다음과 같이 활용한다.
이 그림에서 보면 Loss.backward후 Gradient가 각각 다른 것을 알수 있는데, 이는 All-Gather에서 합쳐진 Weight는 동일하나 각 GPU별로 서로 다른 Mini-Batch가 동작하기 때문에 달라진다. 다시 3의 그림으로 돌아가서 설명하면
Backward는 Unit2의 All-Gather로 시작되며 Unit2에서 Backward가 완료되면 Parameter를 Free하고(동시에 Reduce-Scatter가 진행: RS2) Layer3가 있는 Uni0의 backward가 진행(bwd0), layer2, layer1의 backward를 위해 Unit1의 All-Gather가 진행되고 Unit1의 Backward가 진행되면서 Release-Scatter가 진행 된다. Uni0의 경우 Parameter Release가 최종 단계까지 진행되지 않는다.
따라서 하나의 GPU에서 메모리 요구는 최대 Unit0의 모든 파라미터 + (Unit1 또는 Unit2의 모든 파라미터중 큰 값) + Unit0의 중간 단계 산출물(Gradient, Optimizer State등) + (Unit1 또는 Unit2의 중간 단계 산출물의 최대값)이 된다.
