Keras项目中使用JAX实现多GPU分布式训练指南-优快云博客

Keras项目中使用JAX实现多GPU分布式训练指南

keras keras-team/keras: 是一个基于 Python 的深度学习库，它没有使用数据库。适合用于深度学习任务的开发和实现，特别是对于需要使用 Python 深度学习库的场景。特点是深度学习库、Python、无数据库。项目地址: https://gitcode.com/gh_mirrors/ke/keras

分布式训练概述

在深度学习领域，当模型规模和数据量不断增大时，单设备训练往往会遇到计算资源和内存瓶颈。分布式训练技术应运而生，它主要通过两种方式实现计算能力的扩展：

数据并行：将模型复制到多个设备上，每个设备处理不同的数据批次，最后合并结果。根据同步方式不同，又可分为同步数据并行和异步数据并行。
模型并行：将单个模型的不同部分分配到不同设备上，共同处理同一批次数据。这种方法适合具有天然并行结构的模型。

本文重点介绍同步数据并行方法，这种方法能保持模型收敛行为与单设备训练一致，是最常用的分布式训练方式。

环境准备

在开始之前，我们需要设置JAX作为Keras的后端，并导入必要的库：

import os
os.environ["KERAS_BACKEND"] = "jax"

import jax
import numpy as np
import tensorflow as tf
import keras
from jax.experimental import mesh_utils
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

模型与数据准备

构建简单CNN模型

我们构建一个包含卷积层、批归一化层和Dropout层的简单CNN模型：

def get_model():
    inputs = keras.Input(shape=(28, 28, 1))
    x = keras.layers.Rescaling(1.0 / 255.0)(inputs)
    x = keras.layers.Conv2D(filters=12, kernel_size=3, padding="same", use_bias=False)(x)
    # ... 中间层省略 ...
    outputs = keras.layers.Dense(10)(x)
    return keras.Model(inputs, outputs)

准备MNIST数据集

def get_datasets():
    (x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
    x_train = np.expand_dims(x_train.astype("float32"), -1)
    x_test = np.expand_dims(x_test.astype("float32"), -1)
    return (
        tf.data.Dataset.from_tensor_slices((x_train, y_train)),
        tf.data.Dataset.from_tensor_slices((x_test, y_test))
    )

单主机多设备同步训练实现

核心概念

在这种设置下，一台主机配备多个GPU/TPU（通常2-16个），每个设备运行模型的副本（称为副本）。训练过程中的关键步骤：

数据分片：全局批次被分割为多个本地批次
并行处理：每个副本独立处理本地批次
梯度同步：所有副本的梯度更新在每一步结束时高效合并

实现步骤

配置基本参数：

num_epochs = 2
batch_size = 64
train_data, eval_data = get_datasets()
train_data = train_data.batch(batch_size, drop_remainder=True)

定义计算图和训练步骤：

# 计算损失函数
def compute_loss(trainable_variables, non_trainable_variables, x, y):
    y_pred, updated_non_trainable_variables = model.stateless_call(
        trainable_variables, non_trainable_variables, x
    )
    return loss(y, y_pred), updated_non_trainable_variables

# 计算梯度
compute_gradients = jax.value_and_grad(compute_loss, has_aux=True)

# 训练步骤
@jax.jit
def train_step(train_state, x, y):
    # 解包训练状态
    trainable_variables, non_trainable_variables, optimizer_variables = train_state
    # 计算梯度和损失
    (loss_value, non_trainable_variables), grads = compute_gradients(
        trainable_variables, non_trainable_variables, x, y
    )
    # 更新参数
    trainable_variables, optimizer_variables = optimizer.stateless_apply(
        optimizer_variables, grads, trainable_variables
    )
    return loss_value, (trainable_variables, non_trainable_variables, optimizer_variables)

设备网格与分片策略：

# 获取设备数量
num_devices = len(jax.local_devices())
devices = mesh_utils.create_device_mesh((num_devices,))

# 变量复制策略（所有设备上复制完整变量）
var_mesh = Mesh(devices, axis_names=("_"))
var_replication = NamedSharding(var_mesh, P())

# 数据分片策略（沿批次维度分片）
data_mesh = Mesh(devices, axis_names=("batch",))
data_sharding = NamedSharding(data_mesh, P("batch"))

训练循环：

# 初始化训练状态
train_state = get_replicated_train_state(devices)

for epoch in range(num_epochs):
    for data in iter(train_data):
        x, y = data
        # 数据分片
        sharded_x = jax.device_put(x.numpy(), data_sharding)
        # 执行训练步骤
        loss_value, train_state = train_step(train_state, sharded_x, y.numpy())
    print(f"Epoch {epoch} loss: {loss_value}")

更新模型参数：

训练完成后，需要将分布式训练得到的参数更新回原始模型：

trainable_variables, non_trainable_variables, _ = train_state
for var, value in zip(model.trainable_variables, trainable_variables):
    var.assign(value)
for var, value in zip(model.non_trainable_variables, non_trainable_variables):
    var.assign(value)