다중 GPU 학습

여러 장치에 나누는 두 가지 방법

GPU 한 장으로 부족할 때는 자료나 모델을 여러 장치에 나눌 수 있어. 두 방식은 서로 다른 문제에 답해:

데이터 병렬화 — 모든 GPU에 모델 전체 복사본을 올리고 배치를 서로 다른 조각으로 나눠 처리해. 각 단계에서 장치 사이의 그래디언트를 평균내며, 모델은 한 장치에 들어가지만 학습이 느릴 때 사용해.
모델 병렬화 — 모델 자체가 GPU 한 장에 들어가지 않을 때 가중치를 여러 장치로 나눠. 속도를 논하기 전에 메모리 부족이 나는 경우의 해법이야.

Keras 3의 분산 API

Keras 3는 두 방식을 keras.distribution 아래에 통합했어. 데이터 병렬화는 장치로 DataParallel 객체를 만들고 모델을 구성하기 전에 set_distribution()을 호출해. 그러면 fit()이 배치 분할과 그래디언트 집계를 맡아 학습 코드를 바꿀 필요가 거의 없어. 모델 병렬화에서는 DeviceMesh를 선언하고 각 가중치 텐서를 메시 축에 나누는 방법을 LayoutMap으로 지정해.

단순한 방법부터 적용해

데이터 병렬화는 대략 GPU 8장까지 거의 선형으로 확장되며 추가 코드도 적어. 모델 병렬화는 설계 비용이 크므로 모델 복사본 하나가 실제로 단일 장치에 들어가지 않을 때만 사용해. 먼저 측정하고 선택해.

Code

keras.distribution의 데이터 병렬화와 모델 병렬화·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의 분산 API

단순한 방법부터 적용해

Code

External links

Exercise

Progress

댓글 0