Recent Server Workload

Word Embedding

Embedding : Mapping to a point in some space, 여기서 mapping을 할 때, 그 의미가 있어야 됨

전통적인 word embedding 방법에는 one-hot이 있음

• 굉장히 sparse 함
• one-hot vector의 dimension이 너무 커짐

대신 어떤 dense한 embedding vector를 만들어보자

Word Embedding (Word2Vec)

skip-gram : 주변 단어들을 통해 해당 단어를 예측

large corpus로 train 후 어떤 weight들로 이루어진 벡터로 나타내어짐(embedding vector)

Recommender System on Servers

유튜브 Recsys은 Neural collaborative Filtering을 사용

• user history에 해당하는 user embedding과 item embedding이 주어지고, 딥러닝 layer를 지나 아이템을 추천

이 때 아래와 같은 server측면에서의 문제점들이 있음

Wordload 측면

1. Embedding table을 look-up하는 시간이 너무 많이 듦(main memory access로 인해서) → TLB 필요
2. Fully Connected(FC) layer를 compute하는데도 시간이 너무 많이 듦

Memory 측면

1. Bandwidth
2. Capacity : embedding table이 너무 큼 (수많은 유저와 그것보다 더 많은 영상)

ZION: Facebook Next Generation Large Memory Training Platform

Server Rack에 여러 개의 모듈이 들어감

OCP Accelerator Module (OAM)

이 세상엔 수많은 accelerator가 있고, 각각 다른 system에만 맞추어 제작됨

그래서 accelerator module을 standarize하기 위해 만들어짐

구조

• 8 socket CPU + 8 Accelerator system
• CPU 에는 커다란 embedding table을 저장하고 (DRAM이 사이즈가 GPU Memory보다 훨 크기 때문) MLP Parameter는 GPU main memory에 저장
• MLP를 여러 GPU에 분산하여 저장하므로, Ring network로 각각의 parameter를 공유하도록 함, 이때 all-Reduce Scheme이 사용됨
  • All reduce : 학습 과정에서 각각의 노드에 흩어져 있던 mini-batch의 weight가 업데이트되면 이를 accumulate해서 global weight를 얻는 것
  • All gather : 학습 과정에서 업데이트 된 weight를 copy하는것

ZION 구조를 통해 기존의 방법(GPU Only)보다 memory capacity를 늘릴 수 있었음 (CPU memory를 같이 사용하므로)

Current bottleneck in server

Bandwidth

low bandwidth btw CPU - GPU

SSD in large model training on multi GPU nodes

• NVIDIA Grace CPU를 사용해서 bandwidth를 증가할 수 있었음
• Youtube RS에서 embedding table이 사이즈가 매우 크므로, 이걸 average operation하는걸 GPU나 CPU에서 하지 말고 메인메모리(DIMM)에서 미리 계산해서 넘기면 Bandwidth가 낮아질 수 있다 → TensorDIMM Architecture

Capacity

• DRAM/NVM architecture
  • NVM(non volatile memory)를 사용해서 거대한 embedding table을 저장하도록 함
  • 이때 문제는 low memory BW in Flash Memory
  • basic idea : DRAM Cache를 이용, embedding table access시 나타나는 locality를 활용하자 (마치 main memory의 cache가 되는 역할)
    1. Embedding table data를 partition해서 flash page number를 줄이도록
    2. 자주 사용하는 데이터를 cache에 insert (Adaptive cache insertion)

Self-Supervised Learning

large model은 좋은 성능을 보여줌

Training cost는 지속적으로 증가

Why Self-Supervised Learning?

• labeling is expensive
• server now has abundant useful data

대부분 Self-Supervised Learning model의 작동 방식은, input을 slightly change해서 2개를 만들어 둘의 차이를 없애도록 → input data의 entity를 학습하도록

→ 우리가 원하는 것은 well-trained model, 즉 good representation을 얻는 것

Machine Translation

기존 RNN이나 다른 번역 모델에서는 문장 번역 시 alignment가 문제가 됨 →Transformer의 탄생

Neural Machine Translation (NMT)

• encoder RNN과 decoder RNN을 사용
• but 마지막 Decoder RNN cell에 모든 source sentence의 정보가 담기게 됨 (information bottleneck) → Sequence to Sequence with Attention!

Bert 학습 과정

1. Pretrain with large corpus of data
   1. masked language model : 문장이 주어지면 mask된 단어를 예측하도로 학습
   2. Next Statement Prediction : Auto-Regressive
2. Fine Tuning for Target Task
   1. 각 목적에 맞게

pretrain 후 fine tuning으로 적은 데이터에서 오는 overfitting 문제를 해결 (일반적으로 거대한 model에서 적은 데이터로 학습하면 overfitting 발생하게됨)

Switch Transformer

• Conditional Execution in the 600B Model
  • FFN 여러개를 병렬로 작동시킴 → 다 작동하는 건 아니고, input에 따라 필요한 FFN만 작동시킴
    • 이를 통해 모델 크기는 늘리면서 compute overhead가 낮음

Computation of BERT

• per 1 layer
• H : Hidden Size / L : word 수 / h : head 수
• 각각의 head에서 input과 layer끼리 먼저 곱하는 computation : 3Hdh = 3H^2 (H = hd)
• 각각의 output에 대해 layer를 곱하는 H^2
• FFN에서 4H^2 + 4H^2 = 8H^2
• 따라서 총 12H^2
• 이 때, multihead Attention : Feed Forward Network = 1 : 2
• 즉 FFN이 전체 computation cost를 dominate

Computation of Transformer

• per 1 layer
• H : Hidden Size / L : word 수 / h : head 수
• 각각의 head에서 computation : 3LHd + 2L^2 d
  • 전체 head에 대해서는
  • h(3LHd + 2L^2 d) = 3LH^2 + 2L^2H +LH^2 = 12L^3 + 4L^3 + 4L^3 = 20L^3
• output에 대해 마지막 layer를 곱하는 LH^2
• FFN에서 8LH^2 = 32L^3
• 따라서 총 20L^3 + 32L^3 = 52L^3
• 이 때, multihead Attention : Feed Forward Network = 2 : 3

ChatGPT 학습과정

• chatGPT에서 input은 Prompt(Context)라고 칭함
• Prompt는 request에 대해 best serve를 하도록 디자인됨
• Auto Regressive, 즉 모델 자신이 만들어낸 output을 다시 자신이 학습하는데 사용함

Step 0) Pre-Training

Pretrain with large datasets

Step 1) Supervised Learning with Examples

• Human labeling data나 실존하는 데이터(수능 시험 같은거)로 supervised learning
• 어떤 input에 대해 원하는 output이 나오도록 하기 위함

Step 2) Reward Model Training

• ChapGPT 에 thumbs-up과 같은 score를 통해 Reward를 줄 수 있는 모델을 학습시킴 → Step 3에서 reinforcement learning을 하기 위함
• 즉 Rating을 예측하는 새로운 모델을 만듦

Step 3) Reinforcement Learning using Reward Model

• prompt가 주어지면 ChatGPT는 output을 만들어냄
• Reward Model이 output에 대해 reward를 줌
• Reward를 기반으로 ChatGPT 모델을 fine-tuning

Contrastive Predictive Coding (CPC)

• Target Class를 직접적으로 추정하지 않고 Target 위치의 벡터와 다른 위치의 벡터를 비교하는 방식으로 학습
• 어떤 waveform의 shared information을 통해(그냥 한 signal에서 서로 다른 부분이기에 shared라 하는 것 같음) 보다 global 한 representation을 학습시키도록 (local information과 local noise는 discard, 버리도록)

SimCLR : A Simple Framework for Contrastive Learning of Visual Representations

• input을 transform한 두 개의 다른 input을 만듦
• Resnet-50을 통과시켜 feature map을 얻음
• 2-layer MLP를 통과시켜 zi, zj를 얻음
• 이 zi, zj의 차이를 적게 만들도록 학습시킴 (Maximize agreement)
• 이때, large batch에 대해 학습하면 더 negative sample을 보게 만들 수 있음 → 다른 거는 다르게 학습되고 같은 데이터에 대해서는(positive sample) 같도록 학습, 즉 negative가 많을수록 학습이 잘된다???

Self Supervised Learning으로 label을 줄이면서 더 큰 모델을 효과적으로 학습시킬 수 있었음

BYOL : Bootstrap Your Own Latent

• No negative samples and small batch size

DINO랑 비슷함

Barlow Twins

• 어떤 이미지를 두 tranformation시킴
• 앞에 두 transform된 distorted 이미지를 network를 통과시켜 representation을 얻음
• 둘을 correlation 시킴 (cosine similarity?)
• 이 때 correlation matrix에서 대각선은 1, 나머지는 -1(아예 다르다고)이 되도록 학습시키면 두 distorted image가 같다라는 것을 모델이 학습하게 됨

CLIP이랑 비슷함(Transformer를 사용)

Training ML models at Scale

large activation problem

• large batch training of large activation model은 큰 메모리가 필요, 그런데 GPU memory는 너무 작음
  • 그래서 training 시에만 multiple GPU를 사용
• Adam 과 같은 advanced training 기법을 사용할 경우 원래 모델 size보다 더 큰 메모리가 필요

Activation에 큰 메모리가 필요한 이유

• training 시에는 forward와 backward 연산 시, 특히 back propagation때문에 중간 값들을 저장할 메모리가 필요

해결법

Activation Checkpointing

• forward pass에서 각 layer의 activation을 저장하지 말고, 매 N layer마다의 activation 결과값만 저장
• 이후 저장되지 않은 값은 backward pass시 N-1번째 layer의 결과값을 통해 re-calculate 후 weight update하도록
• computation은 조금 늘어날지라도 가용할 메모리가 늘어났으므로 large model과 large batch size를 적용할 수 있기 때문에 이게 더 이득

Activation Compression

Quantization으로 activation memory cost를 줄일 수 있음

Pooled

CXL 메모리 사용

Data Parallelism

• 하나의 모델의 parameter가 각각의 GPU에 replicated되어 올라감
• mini-batch에 대해 각 GPU가 weight를 계산
• 나중에 all reduce(합치고 평균내서 모든 GPU가 global weight를 가지도록

Model Parallelism

• 하나의 모델의 parameter를 GPU node들이 share
• batch data에 대해 weight를 구한 후 all gather를 통해 모든 GPU가 계산된 하나의 모델에 대한 parameter를 가지도록
• model parallelism은 GPU node 내에서만 가능
  • node - node끼리는 low BW로 traffic이 너무 커짐

Pipeline Parallelism

현재 model parallelism은 large scale training에 적합하지 않음

• GPU node - node사이의 BW가 너무 적고 communication이 느
• Pipeline parallelism으로 극복 가능

Model size보다 더 많은 메모리가 필요

1T parameter model을 train시킬 수 있을까?

• 1T까지는 올릴 수 있지만, 실제로는 더 많은 메모리가 필요 : Adam 과 같은 advanced된 training 기법에서 더 많은 parameter가 필요하기 때문

ZERO(Zero Redundancy Optimizer)

DP MP 모두 전체 모델 parameter를 사용함

그런데 실제로는 전체 training 하는 동안 전체 모델 parameter가 필요없음

ZeRO-DP

• 모델 state(parameter)를 나누고 dynamic communication으로 communication volumn은 줄이면서 본질적인(intrinsic) 모델 state는 얻을 수 있도록
• large model에서는 전체 모델을 저장 못함
• 대신 chunk로 나누고 parameter 연산 후 all gather로 다시 모음, 이를 반복

forward pass

1. 각 GPU node가 모델의 한 layer의 parameter의 chunk를 담당(1번째 chunk는 첫 번째 GPU node가, n번째 chunk는 n번째 node가 담당..)
2. 연산 후 all gather후 output 계산
3. 다음 layer로 넘어갈 때에는 이전 layer의 param들은 discard

backward pass

1. 마지막 layer에서 all gather하고 gradient descent
2. discard하고 위 layer로 이동
3. 반복

All Reduce = ReduceScatter + Allgather

ZeRO Offload

(Adam 같은) optimizer의 state를 GPU가 아닌 CPU에 저장 → GPU memory를 쓸 필요 없음

ZeRO Infinity

CPU, GPU 외에 SSD까지 사용