Keras 3으로 분산 트레이닝하기

import os

# 분산 API는 현재 JAX 백엔드에서만 구현되어 있습니다.
os.environ["KERAS_BACKEND"] = "jax"

import keras
from keras import layers
import jax
import numpy as np
from tensorflow import data as tf_data  # 데이터셋 입력을 위한 모듈.

# 로컬에서 사용 가능한 GPU 장치를 검색합니다.
devices = jax.devices("gpu")  # 8개의 로컬 GPU가 있다고 가정합니다.

# 데이터 및 모델 병렬 축을 가진 2x4 장치 메쉬를 정의합니다.
mesh = keras.distribution.DeviceMesh(
    shape=(2, 4), axis_names=["data", "model"], devices=devices
)

# 텐서가 메쉬에 어떻게 분산되는지 설명하는 2D 레이아웃.
# 레이아웃은 "model"을 행으로, "data"를 열로 하는 2D 그리드로 시각화할 수 있으며,
# 물리적 장치에 매핑될 때 [4, 2] 그리드입니다.
layout_2d = keras.distribution.TensorLayout(axes=("model", "data"), device_mesh=mesh)

# 이미지 입력의 데이터 병렬에 사용할 수 있는 4D 레이아웃.
replicated_layout_4d = keras.distribution.TensorLayout(
    axes=("data", None, None, None), device_mesh=mesh
)

# 장치 목록을 사용하여 DataParallel 생성.
# shortcut으로서 장치 목록을 생략하면,
# Keras가 모든 로컬 가용 장치를 자동으로 감지합니다.
# 예: data_parallel = DataParallel()
data_parallel = keras.distribution.DataParallel(devices=devices)

# 또는 1D `DeviceMesh`로 DataParallel을 생성할 수 있습니다.
mesh_1d = keras.distribution.DeviceMesh(
    shape=(8,), axis_names=["data"], devices=devices
)
data_parallel = keras.distribution.DataParallel(device_mesh=mesh_1d)

inputs = np.random.normal(size=(128, 28, 28, 1))
labels = np.random.normal(size=(128, 10))
dataset = tf_data.Dataset.from_tensor_slices((inputs, labels)).batch(16)

# 글로벌 분산 설정.
keras.distribution.set_distribution(data_parallel)

# 여기서부터 모든 모델 가중치는 `DeviceMesh`의 모든 장치에 복제됩니다.
# 여기에는 RNG 상태, 옵티마이저 상태, 메트릭 등이 포함됩니다.
# `model.fit` 또는 `model.evaluate`에 입력되는 데이터셋은,
# 배치 차원에서 고르게 분할되어 모든 장치로 전달됩니다.
# 수동으로 손실을 집계할 필요가 없습니다.
# 모든 계산은 전역 컨텍스트에서 수행됩니다.
inputs = layers.Input(shape=(28, 28, 1))
y = layers.Flatten()(inputs)
y = layers.Dense(units=200, use_bias=False, activation="relu")(y)
y = layers.Dropout(0.4)(y)
y = layers.Dense(units=10, activation="softmax")(y)
model = keras.Model(inputs=inputs, outputs=y)

model.compile(loss="mse")
model.fit(dataset, epochs=3)
model.evaluate(dataset)

mesh_2d = keras.distribution.DeviceMesh(
    shape=(2, 4), axis_names=["data", "model"], devices=devices
)
layout_map = keras.distribution.LayoutMap(mesh_2d)

# 아래 규칙은 d1/kernel에 매치되는 모든 가중치가 모델 차원(4개의 장치)으로
# 분할된다는 의미입니다. d1/bias도 동일하게 설정됩니다.
# 그 외의 모든 가중치는 완전히 복제됩니다.
layout_map["d1/kernel"] = (None, "model")
layout_map["d1/bias"] = ("model",)

# 레이어 출력의 레이아웃도 설정할 수 있습니다.
layout_map["d2/output"] = ("data", None)

model_parallel = keras.distribution.ModelParallel(layout_map, batch_dim_name="data")

keras.distribution.set_distribution(model_parallel)

inputs = layers.Input(shape=(28, 28, 1))
y = layers.Flatten()(inputs)
y = layers.Dense(units=200, use_bias=False, activation="relu", name="d1")(y)
y = layers.Dropout(0.4)(y)
y = layers.Dense(units=10, activation="softmax", name="d2")(y)
model = keras.Model(inputs=inputs, outputs=y)

# 데이터는 2개의 장치가 있는 "data" 차원으로 분할됩니다.
model.compile(loss="mse")
model.fit(dataset, epochs=3)
model.evaluate(dataset)

full_data_parallel_mesh = keras.distribution.DeviceMesh(
    shape=(8, 1), axis_names=["data", "model"], devices=devices
)
more_data_parallel_mesh = keras.distribution.DeviceMesh(
    shape=(4, 2), axis_names=["data", "model"], devices=devices
)
more_model_parallel_mesh = keras.distribution.DeviceMesh(
    shape=(2, 4), axis_names=["data", "model"], devices=devices
)
full_model_parallel_mesh = keras.distribution.DeviceMesh(
    shape=(1, 8), axis_names=["data", "model"], devices=devices
)