SHPC 21 - AI parallelism

Cluster for bigger AI model

• 큰 모델을 학습시키기 위해 클러스터를 사용하기도 함
  • 통신 오버헤드(트레이닝 스텝별 sync를 위한 오버헤드)가 클 것
  • 오버헤드가 있더라도 큰 모델을 돌릴 수 있는 것이 중요
• 여러 개의 GPU를 사용하는 경우도 마찬가지

Data parallelism

• 동일한 연산을 서로 다른 데이터에 대해 수행하는 것
• 모든 GPU가 모델의 동일한 카피를 가지고 있다
• Procedure
  • 데이터셋을 미니배치로 나눔
    • weight update가 이뤄지는 단위
  • 각 미니배치를 마이크로배치로 나누고, 각 마이크로배치를 서로 다른 GPU에 나눔
  • 각 GPU는 독립적으로 마이크로배치에 대한 gradients 를 계산
    • 각자가 copy한 모델으로 계산
  • all reduce로 해당 gradients를 모음
  • aggregated gradient의 average를 구하고 각각 weight update
• GPU 메모리가 모델 space보다 커야 한다는 단점이 존재 → Model parallelism

Model Parallelism

• 동일한 input에 대해 파라미터를 나눠서 계산하는 것
• 이 또한 통신이 필요
• 각 단계마다 모델 파라미터를 sync할 필요가 있음
  • gradient aggregation
• 일의 granularity가 작아진다
  • GPU의 utilization이 떨어지고, comm overhead가 증가
    • 전자는 하도 모델이 크니 상관업다

Tensor Parallelism and Pipeline Parallelism

• Tensor parallelism
  • 샤딩
  • 모델 parallelism과 비슷하다. 또 용어 만들었다고 투덜대심
  • tensor == weight
  • 나눠진 tensor에 대한 op를 각 GPU에서 수행
• Pipeline parallelism
  • Hidden layer가 여러개. 이 layer 앞 두세개를 GPU 하나에, 뒤는 또 다른 GPU에…
    • Task를 하나씩 파이프라이닝하는 것
  • 서로 다른 GPU가 서로 다른 단계를 동시에 수행

ZeRO (Zero-Redundancy Optimization)

• 중요한 논문
• 아이디어: 모델 파라미터들을 GPU에 distribution하고, 다른 파라미터가 필요하면 필요할 때마다 다른 GPU에서 가져온다.
  • 통신을 늘려서 메모리를 줄이는 것
• 아래 그림의 후자에 해당한다
  • 전자는 data parallelism
• ZeroDP : Data parallelism 기반, 위와 같은 식으로 필요한 파라미터만 가지고 있는 방식
• Zero-Offload: 파라미터 등을 CPU 메인 메모리에 오프로딩하는 방식
• Zero 1, 2, 3, Offload 라는 이름들로 기능이 다양함
• Shared Distributed Data Parallelism
  • Fully Sharded Data Parallelism (FSDP) 이라고도 부름
  • 그냥 이런 이름으로도 부른다고