0. Preliminary: Where Did All the Memory Go?

• 우리는 pytorch 학습시 memory를 어디에서 소모하게 될까요? 모델, 특히 LLM을 학습하게 되면 VRAM이 부족한 상황을 늘 겪게 되는데요. Deepspeed 논문에서 설명하는 VRAM이 소모되는 순간들은 아래과 같습니다.
• (참고로, NCCL operation에 대한 이해도가 있으셔야 글에 대한 이해가 편함으로 [2. List of NCCL Operation ]을 한번 보고 오시기를 권장드립니다!)

Model States: Optimizer States, Gradients and Parameters

• 모델을 학습하기 위해서는 GPU에 크게 3가지가 올라가야 합니다.

  • Model Parameters

  • Gradients

    ⇒ Model Parameters와 Gradients는 동일한 memory footprints를 차지합니다.

  • Optimizers

    ⇒ 어떤 optimizer를 사용하느냐에 따라 크게 달라지는데요, Adam과 같이 gradient의 (1) first order moments(기울기의 지수평균)와 (2) second order moments(기울기 제곱의 지수 평균: 크기가 큰 기울기에는 작은 학습률을, 작은 기울기에는 큰 학습률을 적용:how? 이 term을 분모로 씀)를 동시에 활용하는 optimizer는 Model Parameters의 2배를 memory footprints로 차지시킵니다.

    [참조: Reflections on Optimizer and LM parameter values]

• Model Parameters과 Gradients의 memory footprints를 감소시키기 위해 parameters(config에 따라 다름)와 activations을 half precision으로 쓰는 Mixed Precision이 많이 활용되기도 하는데요.

  • optimizer에서는 (1) first order moments와 (2) second order moments를 계산해야하기 때문에 fp32를 사용합니다.
  • Half Precision으로 Model Copy를 뜨는 순간 실질적인 memory footprints 감소가 크지 않게 되며, Model Copy를 뜨지 않아도 LLM 같이 큰 모델을 학습시 Mixed precision의 효과를 극대화하려면 Half Precision인 activation을 최대한 많이 활용하는 실험환경이어야 된다고 생각합니다. (개인적으로 이거보단 model parameters의 precision를 처음부터 낮추던가 그냥 full precision으로 학습하는게 더 좋은 선택지라고 생각합니다.)

Residual Memory Consumption

• Optimizer States, Gradients and Parameters을 제외하고도 학습중에도 memory footprints를 발생시키는 요인들이 더 있는데요.

  • Activations

    • 가장 대표적인게 hidden state를 통과하고 나온 representations인 activations입니다.
    • 이 activations는 bsz랑 seq_len이 길어짐에 따라 기하급수적으로 memory footprints를 증가시키는데요, 1.5B parameter GPT-2기준 1K, bsz32에서 60GB의 메모리를 발생시킨다고 합니다.
    • Activation memory footprint를 해결하는 방법은 특정 layer의 gradient를 계산할때마다 forwarding을 다시 해주는 Activation checkpointing이나 Gradient Accumulation등을 통해서 Activation으로 인한 과부하를 줄여줄 수 있겠죠.

  • Temporary buffers

    • Sharding을 하다보면 gradient partition에 all-reduce를 호출을 하거나 grad_norm 연산을 하기전에 gradient를 single flattened buffer에 먼저 넣어놓아야하는데요. 이를 위해 allocated 해 놓은 Temporary buffers도 non-trival한 memory를 차지하게 됩니다.

  • Memory Fragmentation

    

    • 아래에서 설명한 다양한 테크닉들이 적용되거나 큰 모델 학습시, 극단적으로 30%의 memory가 남았음에도 할당되지 못한 경우가 존재한다고 합니다.

1. What is Deepspeed?

• Deepspeed는 위에서 설명한 문제점들을 해결하기 위해 code 몇줄만 추가하면 pytorch 위에 wrapping을 진행해 줌으로써 ‘distributed training, mixed precision, gradient accumulation와 같이 모델 개발에 있어서 필요한 기능을 효율적으로 지원해주도폴 개발된 프레임워크라고 생각하시면 됩니다.

• 논문에서 베이스라인으로 두는 방법론은 DP(Data Parallelism)으로 (1) Model, Optimizer를 모두 여러 device에 복사한 후 (2) 각 device에서 다른 Mini-batch로 forward와 backward를 진행한 후 (3) gradient 평균으로 모델을 동기화하는 방법을 언급하고 있습니다.

  • 논문에서 이야기 하는 DP 방법론은 pytorch에서 이야기하는 DP보다는 DDP에 더 가까운 방법론이라고 보는게 맞는거 같습니다.

• 그렇다면 DeepSpeed에서는 대규모 모델의 memory footprint문제를 해결하기 위해서 어떤 테크닉을 제시했을까요?

  • (스포를 하자면) FSDP랑 크게 다르다고 느껴지지는 않으니, FSDP에 익숙하신 분들은 더 빠르게 이해하실 수 있을꺼라고 생각합니다.

2. ZeRO

• DeepSpeed에서는 2개의 optimization을 제안했습니다.
  1. ZeRO-DP: Model States: Optimizer States, Gradients and Parameters 에서 발생한 memory footprint를 감소시키기 위한 방법론을 제안했습니다.
  2. ZeRO-R: Residual Memory Consumption 에서 발생한 memory footprint를 감소시키기 위한 방법론을 제안했습니다.