Multi-GPU 학습

모델 퍼뜨리는 두 가지 방법

GPU 한 개로 부족하면 *뭔가* 를 여러 개에 나눠. 나눌 수 있는 게 두 가지인데, 각각 다른 부족함에 답해:

data parallelism — 모든 GPU 에 model 전체 사본 올리고 batch 를 GPU 별로 다른 slice 먹여. gradient 는 매 step device 간 평균. *모델은 들어가는데 학습이 느릴 때* 의 답.
model parallelism — model 자체가 GPU 한 개에 안 들어가서 *weight* 를 device 사이 분할. *느려지기도 전에 out-of-memory* 부터 나는 경우의 답.

Keras 3 distribution API

Keras 3 는 둘 다 keras.distribution 아래로 통합했어. data parallel 은 device 로 DataParallel 객체 만들고 model 만들기 *전* 에 set_distribution() 호출 — 그러면 fit() 이 batch sharding 이랑 gradient reduce 를 알아서 해줘, 학습 코드 한 줄 안 바꾸고. model parallel 은 더 의도적이야 — DeviceMesh 선언하고, 각 weight tensor 가 mesh 축에 어떻게 sharding 되는지 적는 LayoutMap 을 줘.

쉬운 거 먼저 잡아

data parallel 은 대략 8 GPU 까지 거의 선형으로 scale 하고 추가 코드가 거의 없어. model parallel 은 손 많이 가고 replica 한 개가 진짜 안 들어갈 때만 정당화돼 — 그러니 잡기 전에 측정해.

Code

keras.distribution 으로 data parallel vs model parallel·python

# Data parallelism: same model on each GPU, split data
devices = keras.distribution.list_devices("gpu")
data_parallel = keras.distribution.DataParallel(devices=devices)

# Set distribution before building the model
keras.distribution.set_distribution(data_parallel)

model = build_model()
model.compile(optimizer="adam", loss="mse")
model.fit(x_train, y_train)  # Automatically distributed!

# Model parallelism: split model across GPUs
device_mesh = keras.distribution.DeviceMesh(
    shape=(2,), axis_names=["model"], devices=devices
)
layout_map = keras.distribution.LayoutMap(device_mesh)
layout_map["dense/kernel"] = keras.distribution.TensorLayout(["model", None])

모델 퍼뜨리는 두 가지 방법

Keras 3 distribution API

쉬운 거 먼저 잡아

Code

External links

Exercise

Progress

댓글 0