Keras 3 架构解析（七）：高性能计算系统——从混合精度、JIT 编译到分布式训练与内存优化

引言

在前六篇文章中，我们分别介绍了 Keras 3 的基础计算层、神经网络层、模型构建层、训练系统、数据处理系统以及训练控制与监控系统。

这些部分共同构成了一套完整的深度学习开发体系：数据处理系统负责准备训练样本，Model 与 Layer 负责定义神经网络结构，Keras Ops 提供统一的数学计算接口，训练系统通过自动微分与优化器更新模型参数，Callback 则负责管理训练生命周期。

然而，能够正确执行模型训练，与能够高效执行模型训练，是两个不同的问题。

随着神经网络规模不断扩大，训练系统需要处理的计算量、参数数量和中间激活数据也随之增加。对于大型 Transformer、卷积网络和长序列模型而言，单次训练可能涉及数十亿次浮点运算，并需要在 CPU、GPU 或 TPU 之间传输大量数据。

在这种情况下，训练效率不再仅仅取决于模型的数学结构，还受到数值精度、计算图编译、张量形状、算子实现、内存访问、设备通信和训练调度等因素的共同影响。

例如，两个数学结构完全相同的模型，即使使用相同的计算设备，也可能因为不同的数值精度或编译策略，表现出显著不同的训练吞吐量。

同样，增加 Batch Size 或计算设备数量，也不必然使训练速度按比例提高。

Keras 3 的高性能计算体系，主要通过统一的框架接口与底层计算后端协作，实现这些性能优化能力。

其核心组成包括：

1. 数值精度管理：通过 FP32、FP16、BF16 及混合精度策略平衡计算性能与数值稳定性。
2. 计算图编译：通过 JIT 编译、算子融合及后端编译器降低执行开销。
3. 分布式计算：通过数据并行、模型并行及设备网格将计算扩展到多个加速器。
4. 训练批次优化：通过 Batch Size、梯度累积和执行步骤组合调整计算粒度。
5. 内存优化：通过混合精度、重计算和张量分片降低内存压力。
6. 性能分析：通过训练吞吐量、设备利用率、计算效率和 Profiler 定位瓶颈。

需要明确的是，Keras 并不独立实现所有这些技术。

它提供统一的模型、训练和部分性能优化接口，而实际的编译器、自动微分引擎、设备通信机制和硬件内核，仍主要由 JAX、TensorFlow 或 PyTorch 等后端负责。

本文将从计算原理和软件架构两个角度出发，介绍 Keras 3 的高性能计算系统，重点分析不同优化技术的作用机制、适用条件及其之间的关系。

---

一、高性能训练的基本目标

1.1 训练性能由哪些因素决定

一个深度学习训练步骤通常包含：

[
\text{Forward}
\rightarrow
\text{Loss}
\rightarrow
\text{Backward}
\rightarrow
\text{Optimizer Update}
]

但实际系统中的训练耗时，并不完全由这些数学计算决定。

一次训练步骤还可能涉及：

• 数据准备和传输；
• Python 与框架运行时调度；
• 计算图执行；
• GPU 或 TPU 内核启动；
• 设备内存读写；
• 多设备通信；
• 指标统计与日志输出。

因此，可以从概念上将单步训练耗时表示为：

[
T_{\text{step}}

T_{\text{compute}}
+
T_{\text{memory}}
+
T_{\text{communication}}
+
T_{\text{overhead}}
]

其中：

• (T_{\text{compute}})：数值计算耗时；
• (T_{\text{memory}})：数据访问与内存传输耗时；
• (T_{\text{communication}})：设备之间的通信耗时；
• (T_{\text{overhead}})：框架调度、编译执行管理及其他开销。

这一表达式只是一种逻辑分解。

在实际系统中，部分计算、数据传输和通信过程可以同时进行，因此各项耗时并不总能直接相加。

高性能训练的目标，就是在满足模型数值正确性和训练要求的前提下，尽可能减少这些开销对整体执行时间的影响。

1.2 训练性能的主要评价指标

评价深度学习系统的性能，通常需要同时考虑多个指标。

指标	含义	主要用途
Step Time	单个训练步骤的执行时间	衡量执行延迟
Samples/s	每秒处理的训练样本数量	衡量训练吞吐量
FLOP/s	每秒实际执行的浮点运算量	衡量计算吞吐量
MFU	模型计算相对于硬件理论峰值的利用程度	衡量模型计算效率
Peak Memory	峰值设备内存占用	衡量内存需求
Compilation Time	计算图编译时间	衡量编译成本
Communication Time	多设备通信耗时	衡量分布式扩展成本

其中，Samples/s 通常是比较训练配置时最直接的吞吐量指标。

但仅观察 Samples/s 仍然不够。

例如，增大 Batch Size 可能提高每秒处理的样本数量，却减少相同时间内执行的参数更新次数。

对于固定训练步数与固定训练样本数量两种实验目标，最终训练时间和收敛表现可能不同。

因此，高性能优化必须区分两个目标：

系统性能优化：提高单位时间内完成的有效计算量。

训练算法优化：提高达到目标模型质量的速度。

两者密切相关，但并不完全等价。

1.3 计算密集型与内存密集型

不同神经网络操作具有不同的计算特征。

例如，大规模矩阵乘法通常具有较高的算术强度（Arithmetic Intensity），即相对于内存访问量，需要执行较多浮点运算。

这类操作更容易充分利用现代 GPU 的 Tensor Core 或 TPU 的矩阵计算单元。

相反，逐元素激活函数、部分归一化操作和简单张量变换，可能需要读取与写入大量数据，但实际执行的浮点运算较少。

这类操作可能更容易受到内存带宽限制。

从 Roofline 性能模型出发，运算的可达到性能可以概念性地表示为：

[
P_{\text{attainable}}
\leq
\min
\left(
P_{\text{peak}},
I\times B_{\text{memory}}
\right)
]

其中：

• (P_{\text{peak}})：硬件理论计算峰值；
• (I)：算术强度；
• (B_{\text{memory}})：内存带宽。

这一关系说明，提高硬件理论 FLOP/s 并不必然使所有操作同比例加速。

如果模型主要受到内存带宽或运行时调度限制，那么单纯增加计算能力可能无法显著改善整体吞吐量。

---

二、数值精度：FP32、FP16 与 BF16

2.1 为什么数值精度影响性能

深度学习模型使用浮点数表示参数、激活值与梯度。

不同浮点格式具有不同的存储位宽、指数范围和有效精度。

常见格式包括：

格式	总位数	指数位	尾数位	主要特点
FP32	32	8	23	数值精度较高
FP16	16	5	10	精度较高于 BF16，但动态范围较小
BF16	16	8	7	动态范围接近 FP32，但尾数精度较低

上述位数不包括隐含的有效数字位。

从存储角度看，FP16 和 BF16 都只需要 2 字节，而 FP32 需要 4 字节。

因此，对于相同数量的元素：

[
M_{\text{FP16}}

M_{\text{BF16}}
\approx
\frac{1}{2}M_{\text{FP32}}
]

但这只是单个张量的存储关系。

实际模型训练的总内存还包括权重、梯度、优化器状态以及不同数据类型的中间激活值，因此不能简单认为启用 FP16 后，训练内存一定减少 50%。

2.2 FP16 的数值特征

FP16 具有较多的尾数位，但指数范围相对有限。

这一特征意味着，FP16 在能够表示的数值范围内具有一定的相对精度优势，但更容易发生上溢和下溢。

例如，在深度学习训练中，某些梯度可能具有非常小的绝对值。

如果直接转换为 FP16，这些梯度可能被舍入为零。

这会导致部分参数无法获得预期的更新。

因此，使用 FP16 训练时，需要特别关注梯度的数值范围。

2.3 BF16 的数值特征

BF16 保留了与 FP32 相同宽度的指数位，因此具有接近 FP32 的动态范围。

但 BF16 只有 7 个显式尾数位，因此其有效数值精度低于 FP16。

这一设计使 BF16 在训练过程中通常具有更好的动态范围特性，减少了因数值过小或过大而发生溢出的风险。

但 BF16 并不意味着没有数值误差。

对于对细小数值差异特别敏感的操作，BF16 的较低尾数精度仍然可能产生明显影响。

因此，BF16 更适合被理解为一种能够兼顾计算性能与动态范围的低精度格式，而不是 FP32 的无条件替代品。

2.4 不同硬件上的精度选择

现代 GPU 和 TPU 通常包含针对低精度矩阵运算优化的硬件单元。

因此，在支持相应数值格式的设备上，低精度矩阵乘法和卷积通常能够获得较高的计算吞吐量。

常见选择可以概括为：

硬件场景	常见精度策略
CPU 开发与调试	FP32
支持 FP16 的 NVIDIA GPU	Mixed FP16
支持 BF16 的现代 NVIDIA GPU	Mixed BF16
Google TPU	Mixed BF16
对数值稳定性要求较高的运算	根据需要保留 FP32

这里并不存在适用于所有模型的绝对最优选择。

例如，某些现代 NVIDIA GPU 同时具备较高的 FP16 和 BF16 矩阵计算性能，最终选择可能更多取决于模型的数值稳定性与算子支持情况。

此外，硬件支持某种数据类型，并不代表所有算子都能够以该类型获得相同的加速效果。

---

三、Keras 的混合精度管理

3.1 为什么需要混合精度

如果将整个神经网络的所有变量和运算都转换为 FP16 或 BF16，可能降低部分计算的数值稳定性。

尤其是梯度累积、归约运算和优化器状态更新等操作，通常需要更高的数值精度。

因此，现代深度学习框架通常采用混合精度训练（Mixed Precision Training）。

其基本思想是：

在适合低精度执行的部分使用 FP16 或 BF16，而在需要更高精度的部分保留 FP32。

例如，某个 Dense 层可以使用 BF16 执行主要前向计算，同时使用 FP32 保存模型权重。

这种方式有助于降低部分计算和激活数据的存储成本，同时避免将所有模型状态都限制在低精度格式中。

3.2 DTypePolicy

Keras 3 通过 DTypePolicy 管理层的计算精度和变量精度。

可以将其理解为两个相互关联的属性：

[
\text{DTypePolicy}

(\text{Compute DType},\text{Variable DType})
]

其中：

• compute_dtype：主要计算使用的数据类型；
• variable_dtype：变量存储使用的数据类型。

例如：

import keras

keras.config.set_dtype_policy(
    "mixed_bfloat16"
)

layer = keras.layers.Dense(128)

print(layer.compute_dtype)
print(layer.variable_dtype)

通常得到：

bfloat16
float32

这意味着该层的计算采用 BF16 策略，而变量以 FP32 保存。

Keras 官方的混合精度策略主要包括：

策略	计算类型	变量类型
float32	FP32	FP32
mixed_float16	FP16	FP32
mixed_bfloat16	BF16	FP32

需要注意，compute_dtype 并不意味着该层内部所有算子都必须以相同精度执行。

具体运算仍可能根据数值稳定性、算子实现和后端策略使用不同的内部精度。

3.3 全局精度策略

Keras 允许设置全局默认策略：

keras.config.set_dtype_policy(
    "mixed_float16"
)

或者：

keras.config.set_dtype_policy(
    "mixed_bfloat16"
)

通常应当在创建模型层之前设置全局精度策略。

这样，后续创建的 Layer 会默认采用指定策略。

例如，对于 TPU 环境：

import os

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

import keras

keras.config.set_dtype_policy(
    "mixed_bfloat16"
)

然后创建模型：

model = keras.Sequential([
    keras.Input(shape=(64,)),

    keras.layers.Dense(
        256,
        activation="gelu"
    ),

    keras.layers.Dense(
        128,
        activation="gelu"
    ),

    keras.layers.Dense(1)
])

此时，模型中的普通计算层会采用全局混合精度策略。

3.4 为特定层保留 FP32

Keras 允许为单个 Layer 覆盖默认精度。

例如：

output = keras.layers.Dense(
    1,
    dtype="float32"
)(x)

即使全局设置为 mixed_bfloat16，该输出层仍采用 FP32 策略。

这种方式在某些回归模型中具有实际意义。

例如，模型主体使用 BF16 执行矩阵计算，而最后的输出层使用 FP32，以降低最终预测输出的舍入误差。

不过，设置最后一层为 FP32，并不意味着整个模型的损失计算和梯度都自动具有 FP32 精度。

对于复杂损失函数和特殊数学操作，仍需要检查实际计算路径。

---

四、Loss Scaling：FP16 训练的数值稳定性

4.1 梯度下溢问题

在神经网络反向传播过程中，某些梯度可能非常小。

假设某个参数的梯度为：

[
g=10^{-8}
]

如果直接使用 FP16 表示，这个数值可能小于相应运算能够可靠保留的范围。

结果可能是梯度被舍入为零。

如果大量梯度出现类似情况，就会影响模型参数更新。

4.2 损失缩放的基本思想

Loss Scaling 通过在计算梯度之前，将损失函数乘以一个缩放系数 (S)，提高中间梯度的数值尺度。

设原始损失为：

[
L(\theta)
]

定义缩放损失：

[
L’(\theta)=S L(\theta)
]

则：

[
\nabla_\theta L’

S\nabla_\theta L
]

如果能够正确完成梯度计算，就可以在参数更新之前将梯度除以 (S)，恢复原来的梯度尺度：

[
\nabla_\theta L

\frac{1}{S}\nabla_\theta L’
]

这种方法能够降低部分梯度在低精度计算过程中发生下溢的风险。

但缩放系数过大，也可能导致数值溢出。

因此，实际训练通常使用动态损失缩放，根据梯度的有限性调整缩放系数。

4.3 Keras 的自动 Loss Scaling

Keras 3 的 compile() 提供：

auto_scale_loss=True

当模型采用：

"mixed_float16"

策略时，如果启用自动损失缩放，Keras 可以自动将优化器包装为 LossScaleOptimizer。

例如：

keras.config.set_dtype_policy(
    "mixed_float16"
)

model.compile(
    optimizer=keras.optimizers.Adam(
        learning_rate=1e-3
    ),
    loss="mse",
    auto_scale_loss=True
)

在使用内置训练循环时，Keras 会管理相应的损失缩放机制。

对于 BF16，通常不需要采用与 FP16 相同的动态损失缩放策略，因为 BF16 的指数范围更大。

但需要强调，Loss Scaling 主要用于解决梯度动态范围问题，并不能消除所有低精度数值误差。

---

五、JIT 编译：从 Python 程序到优化后的计算图

5.1 Python 调度开销

在深度学习训练中，一次前向传播可能包含大量数学操作。

如果每个操作都需要由 Python 分别调度，那么即使底层 GPU 或 TPU 计算很快，程序仍然可能受到 Python 解释器和框架调用开销的限制。

对于计算规模较小、操作数量较多的模型，这种问题尤其明显。

例如：

x = ops.matmul(x, w1)

x = ops.relu(x)

x = ops.matmul(x, w2)

x = ops.add(x, b)

这些操作在数学上可以构成一个连续的计算过程。

如果底层编译器能够识别整个计算结构，就可能将部分操作融合，减少中间数据传输和运行时调度开销。

这正是 JIT 编译（Just-In-Time Compilation）的重要用途。

5.2 编译执行的基本过程

JIT 编译通常包含以下阶段：

          Python / Keras
                │
                ▼
          计算图追踪
                │
                ▼
          中间表示 IR
                │
                ▼
          编译器优化
                │
         ┌──────┴──────┐
         │             │
         ▼             ▼
      算子融合      内存与布局优化
         │             │
         └──────┬──────┘
                ▼
          设备执行程序
                │
                ▼
          GPU / TPU 执行

编译器可能执行以下优化：

• 合并相邻的计算操作；
• 消除无用计算；
• 调整中间张量的内存布局；
• 减少临时张量的读写；
• 优化矩阵运算的执行方式；
• 生成适合目标设备的计算程序。

需要注意，这些优化是否能够实际发生，取决于模型结构、算子类型以及后端编译器的支持情况。

JIT 并不能保证所有模型都获得加速。

5.3 不同后端的编译机制

Keras 的多后端设计并不意味着所有后端共用同一个编译器。

Keras 后端	主要编译机制	典型特点
JAX	JIT、XLA	函数追踪与编译优化
TensorFlow	tf.function、XLA	图执行与编译优化
PyTorch	torch.compile	图捕获、图变换及后端代码生成

这些系统都可以减少部分 Python 调度开销，但实现方式和支持范围有所不同。

例如，JAX 通常需要将计算函数转换为适合函数式执行和编译的形式。

PyTorch 则通过图捕获等机制，将符合条件的 Python 程序转换为可优化的计算图。

5.4 Keras 的 jit_compile

Keras 通过 Model.compile() 提供统一入口：

model.compile(
    optimizer="adam",
    loss="mse",
    jit_compile="auto"
)

根据 Keras 3 官方 API：

• JAX 后端在模型支持的情况下使用 XLA 编译；
• TensorFlow 后端在模型支持的情况下使用 XLA 编译；
• PyTorch 后端在 "auto" 模式下默认采用即时执行；
• 对 PyTorch 显式设置 jit_compile=True，可以使用基于 torch.compile 的编译路径。

也可以显式启用：

model.compile(
    optimizer="adam",
    loss="mse",
    jit_compile=True
)

但如果模型包含不受支持的操作，编译过程可能失败。

因此，jit_compile="auto" 通常适合作为初始配置，而是否显式启用编译，应通过实际测试决定。

5.5 编译成本与稳定运行性能

JIT 编译并不是没有成本的。

训练任务首次执行时，可能需要进行：

[
T_{\text{first}}

T_{\text{trace}}
+
T_{\text{compile}}
+
T_{\text{execute}}
]

后续执行相同计算结构时，如果可以复用已编译程序，则通常不再需要支付完整编译成本。

因此：

[
T_{\text{steady}}
\approx
T_{\text{execute}}
]

这意味着首次训练步骤的耗时可能显著高于后续步骤。

对于需要反复执行大量训练步骤的任务，编译成本可能被长期运行收益摊薄。

但对于只运行少量步骤的小型任务，JIT 编译未必具有整体时间优势。

因此，性能评估应当分别记录：

1. 首次编译耗时；
2. 稳定执行阶段的吞吐量；
3. 整个训练任务的总耗时。

不能仅根据首次执行时间判断模型训练效率。

5.6 动态形状与重复编译

编译器通常需要根据输入张量的形状生成相应的计算程序。

如果输入形状频繁发生变化，就可能触发额外的计算图追踪或编译。

例如：

Batch A: (64, 128, 32)
Batch B: (64, 256, 32)
Batch C: (64, 512, 32)

这些输入具有不同的序列长度。

在某些编译配置下，它们可能对应不同的编译结果。

如果模型需要处理大量不同长度的序列，就可能产生较高的编译开销。

因此，在 JIT 编译环境中，通常需要考虑：

• 固定 Batch Size；
• 将可变长度序列填充到有限的长度集合；
• 尽量减少不必要的形状变化；
• 复用能够共享的编译结果。

但需要注意，固定形状可能增加 Padding 带来的无效计算。

因此，减少编译次数与降低计算浪费之间，也存在需要权衡的关系。

5.7 JAX 的持久化编译缓存

对于使用 JAX 后端的任务，可以设置持久化编译缓存目录。

例如：

import jax

jax.config.update(
    "jax_compilation_cache_dir",
    "/tmp/jax_compile_cache"
)

应当在首次执行相关编译操作之前完成配置。

持久化缓存允许后续进程在满足缓存复用条件时使用此前生成的编译结果。

但缓存是否命中仍然取决于计算图、输入形状、设备和编译配置等因素。

因此，不能假定只要使用相同的 Python 模型类，就一定能够复用所有编译结果。

---

六、steps_per_execution：降低训练调度成本

6.1 为什么需要组合多个训练步骤

即使模型已经完成 JIT 编译，训练循环仍然可能需要频繁调用已编译函数。

例如：

Python
  │
  ▼
执行 Batch 1
  │
  ▼
返回 Python
  │
  ▼
执行 Batch 2
  │
  ▼
返回 Python

对于单个 Batch 计算量较小的模型，频繁进入和退出已编译函数可能产生一定的调度成本。

Keras 通过 steps_per_execution 减少这种开销。

6.2 基本使用

例如：

model.compile(
    optimizer="adam",
    loss="mse",
    jit_compile="auto",
    steps_per_execution=16
)

这表示允许一次编译函数执行最多 16 个训练 Batch。

其概念性执行关系为：

             Python
                │
                ▼
         Compiled Function
                │
       ┌────────┼────────┐
       ▼        ▼        ▼
    Batch 1  Batch 2  ... Batch 16
       │        │        │
       └────────┼────────┘
                ▼
             Python

通过减少 Python 调度频率，某些模型可以获得更高的训练吞吐量。

6.3 与 Batch Size 的区别

steps_per_execution 与 Batch Size 不是同一个概念。

假设：

batch_size = 64
steps_per_execution = 16

如果没有梯度累积，那么模型仍然通常对每个包含 64 个样本的 Batch 分别计算梯度并更新参数。

只不过，多个训练 Batch 可以被组织在同一次编译函数执行中。

因此：

[
\text{steps_per_execution}
\neq
\text{gradient_accumulation_steps}
]

也不等于一次梯度更新使用的样本数量。

6.4 适用条件与限制

steps_per_execution 可能适合：

• TPU 训练；
• 单步计算时间较短的模型；
• Python 调度开销较大的任务；
• 具有较稳定输入形状的训练流程。

但需要注意两个限制。

首先，根据 Keras 3 官方文档，该参数不支持 PyTorch 后端。

其次，当一个编译执行包含多个 Batch 时，Batch 级 Callback 的调用频率也会改变。

因此，若需要在每一个 Batch 后记录精确日志，就需要考虑这一参数与监控机制之间的关系。

对于高吞吐量训练，应通过实验比较不同 steps_per_execution 的实际收益，而不是无条件设置较大的数值。

---

七、分布式训练：从单设备到多设备

7.1 为什么需要分布式训练

单个 GPU 或 TPU 的计算能力与设备内存都是有限的。

随着模型规模或训练数据量增加，可能出现两类问题。

第一类是计算时间过长。

即使模型能够放入单个设备内存，训练过程仍可能需要大量时间。

第二类是模型无法放入单个设备内存。

例如，大型模型的权重、梯度、优化器状态和激活值，可能超过单个加速器的容量。

分布式训练主要通过两类方式解决这些问题：

数据并行（Data Parallelism）

在多个设备上复制相同的模型，并让不同设备处理不同的数据。

模型并行（Model Parallelism）

将模型的参数或计算过程分配到多个设备上。

两者可以独立使用，也可以组合使用。

7.2 数据并行的基本原理

假设有 4 个 GPU。

数据并行可以为每个 GPU 分配一份相同的模型。

每个 GPU 处理不同的输入批次：

                Global Batch
                     │
        ┌────────────┼────────────┐
        ▼            ▼            ▼
      Batch 1      Batch 2      Batch 3 ... 
        │            │            │
        ▼            ▼            ▼
       GPU 0        GPU 1        GPU 2 ...
        │            │            │
        ▼            ▼            ▼
      Model A      Model A      Model A
        │            │            │
        ▼            ▼            ▼
      Gradient     Gradient     Gradient
        │            │            │
        └────────────┼────────────┘
                     ▼
              Gradient Sync
                     │
                     ▼
              Update Weights

模型副本具有相同的参数结构。

各设备分别计算当前数据分片对应的梯度，再通过同步机制聚合梯度。

对于采用相同权重、相同损失归约规则的同步数据并行训练，可以将最终梯度理解为不同设备局部梯度的加权聚合。

在局部 Batch 大小相同的常见情况下，通常可以采用平均形式：

[
g=
\frac{1}{D}
\sum_{d=1}^{D}g_d
]

其中，(D) 为设备数量。

7.3 数据并行的性能限制

在理想情况下，增加设备数量可以提高训练吞吐量。

但真实系统还需要考虑设备通信成本。

假设单设备训练吞吐量为：

[
P_1
]

使用 (D) 个设备后的吞吐量为：

[
P_D
]

并行效率可以定义为：

[
E_D=
\frac{P_D}{D P_1}
]

如果：

[
E_D=1
]

表示实现理想线性扩展。

但在实际环境中，梯度同步、设备等待、负载不均衡和数据供应等问题可能降低并行效率。

例如，如果模型非常小，那么每次执行的计算量可能不足以覆盖梯度同步成本。

此时，增加设备数量不一定提高整体性能。

7.4 Keras 的 DataParallel API

Keras 3 提供：

keras.distribution.DataParallel

例如，在 JAX 后端中：

import os

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

import keras

devices = keras.distribution.list_devices()

distribution = keras.distribution.DataParallel(
    devices=devices
)

keras.distribution.set_distribution(
    distribution
)

随后创建模型并进行训练：

model = keras.Sequential([
    keras.Input(shape=(64,)),
    keras.layers.Dense(128, activation="relu"),
    keras.layers.Dense(1)
])

model.compile(
    optimizer="adam",
    loss="mse"
)

在这种配置下，Keras 会利用底层分布式机制，在设备之间分配输入数据并组织计算。

需要特别注意：

Keras 3 官方通用 keras.distribution API 的文档目前主要以 JAX 后端实现为基础。

TensorFlow 和 PyTorch 虽然也支持多设备训练，但不能据此认为同一套 keras.distribution 代码已经在全部后端中具有完全一致的实现。

对于 TensorFlow，可以使用 tf.distribute 等原生分布式机制。

对于 PyTorch，可以使用 DistributedDataParallel 等原生分布式机制。

因此，需要区分 Keras 的统一分布式 API 与不同后端自身的分布式训练能力。

---

八、ModelParallel：模型分片与设备网格

8.1 模型并行的基本原理

数据并行要求每个设备保存相同的模型副本。

因此，它不能直接解决单个模型无法放入单个设备内存的问题。

模型并行则通过将模型参数或计算过程分配到多个设备上，降低单设备的资源需求。

例如，某个较大的权重矩阵为：

[
W\in\mathbb{R}^{4096\times8192}
]

如果沿输出特征维度分成 4 份，则每个分片可以表示为：

[
W_i\in\mathbb{R}^{4096\times2048}
]

不同设备分别保存不同的参数分片。

计算时，再根据具体分片方式组织局部矩阵运算与必要的通信。

这种方式可以降低单个设备需要保存的参数数量。

但它也可能增加设备之间的通信量。

因此，模型并行不是简单地将权重平均分配到多个设备，而是需要根据计算依赖关系设计合理的分片布局。

8.2 DeviceMesh

Keras 通过 DeviceMesh 将物理设备组织为逻辑设备网格。

例如，假设有 8 个设备，可以构建：

[
2\times4
]

的设备网格。

其中：

• 第一个维度用于数据并行；
• 第二个维度用于模型并行。

对应的概念结构为：

                 Model Axis
               M0  M1  M2  M3

Data Axis  D0   ●   ●   ●   ●
           D1   ●   ●   ●   ●

Keras 示例：

devices = keras.distribution.list_devices()

mesh = keras.distribution.DeviceMesh(
    shape=(2, 4),
    axis_names=("data", "model"),
    devices=devices
)

该示例要求实际可用设备数量与网格形状匹配，即共有 8 个设备。

8.3 TensorLayout 与 LayoutMap

仅定义设备网格还不够。

开发者还需要指定张量的不同维度如何映射到设备网格。

Keras 提供两个主要概念：

TensorLayout

描述张量维度与设备网格维度之间的对应关系。

LayoutMap

根据模型变量路径，为不同权重或张量分配相应的布局。

例如：

layout_map = keras.distribution.LayoutMap(
    mesh
)

layout_map[".*projection/kernel"] = (
    None,
    "model"
)

layout_map[".*projection/bias"] = (
    "model",
)

这表示对匹配规则的 Dense 权重和偏置，沿设备网格的 model 轴进行分片。

其中，None 表示该张量维度不沿对应设备网格轴进行分片。

随后可以创建：

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

并将其设置为当前分布式配置。

在实际项目中，需要检查模型变量的完整路径，并确认 LayoutMap 规则与对应张量形状兼容。

此外，所有变量都不一定适合采用相同的分片方式。

例如，某些层的输入和输出张量可能需要额外的跨设备通信，简单地分片权重反而可能降低性能。

8.4 模型并行与通信成本

模型并行可以降低单设备内存需求，但它通常需要在设备之间交换中间结果。

例如，在某些张量并行布局中，多个设备分别计算输出的不同部分。

后续计算可能需要通过 All-Gather、All-Reduce 或其他集体通信操作组合这些结果。

因此，模型并行的性能与设备间互连带宽和通信延迟密切相关。

如果模型参数能够轻松放入单个设备，且计算任务规模较小，那么不必要的模型分片可能降低整体效率。

通常只有在单设备内存不足，或者模型结构能够有效利用并行计算时，模型并行才更有意义。

---

九、Batch Size、Micro-batch 与梯度累积

9.1 Batch Size 的三种含义

在分布式训练中，Batch Size 可能具有不同的含义。

Local Batch Size

单个设备在一次训练步骤中处理的样本数量。

Global Batch Size

多个数据并行设备在一次同步训练步骤中处理的样本总数。

Effective Batch Size

考虑梯度累积后，一次参数更新对应的总样本数量。

假设：

• 每个设备处理 32 个样本；
• 使用 8 个数据并行设备；
• 每累积 4 个训练批次更新一次参数。

那么：

[
B_{\text{global}}=32\times8=256
]

有效批大小为：

[
B_{\text{effective}}

32\times8\times4

1024
]

一般可以写为：

[
B_{\text{effective}}

B_{\text{local}}
\times D
\times A
]

其中：

• (D)：数据并行设备数量；
• (A)：梯度累积次数。

需要注意，在 Keras 内置 fit() 与 DataParallel 配合使用时，传入的 Batch Size 通常表示全局 Batch，而不是每个设备的局部 Batch。

因此，在编写分布式训练配置时，必须明确 Batch Size 的具体语义。

9.2 梯度累积的基本原理

假设一个设备由于内存限制，只能一次处理 32 个样本。

但训练算法希望使用 128 个样本计算一次参数更新。

可以连续处理 4 个 Micro-batch：

[
B_1,B_2,B_3,B_4
]

分别计算梯度：

[
g_1,g_2,g_3,g_4
]

然后取平均：

[
g=
\frac{1}{4}
\sum_{i=1}^{4}g_i
]

最后执行一次优化器更新：

[
\theta_{t+1}

\operatorname{OptimizerUpdate}
(\theta_t,g)
]

这种方法称为梯度累积（Gradient Accumulation）。

9.3 Keras 的梯度累积

Keras 3 的 Optimizer API 提供：

gradient_accumulation_steps

例如：

optimizer = keras.optimizers.AdamW(
    learning_rate=1e-3,
    weight_decay=1e-4,
    gradient_accumulation_steps=4
)

随后：

model.compile(
    optimizer=optimizer,
    loss="mse"
)

在这种配置下，优化器不会在每个训练批次后立即更新参数，而是按照配置累积多个批次的梯度。

根据 Keras 官方 Optimizer 文档，累积期间使用梯度平均值进行参数更新。

需要注意，Optimizer 内部的迭代计数、EMA 更新频率与学习率调度，可能采用不同的计数语义。

因此，不能无条件把数据处理批次数量等同于实际参数更新次数。

9.4 梯度累积与大 Batch 是否等价

在理想条件下，如果每个 Micro-batch 的损失归约方式一致，模型在累积期间不更新参数，那么梯度平均可以与直接使用同等样本数量的大 Batch 具有相同的数学目标。

但是，实际训练中不一定完全等价。

例如：

• Batch Normalization 使用当前批次的统计量；
• Dropout 等随机操作可能采用不同的随机状态；
• 不同批次大小可能导致不同的数值计算顺序；
• 损失归约与样本权重处理可能存在差异；
• 不完整批次可能影响梯度平均方式。

因此，梯度累积更准确的定位是：

在限制单次激活内存的同时，近似实现较大有效批大小的优化方式。

它主要改变参数更新频率与内存需求，不应被直接视为提高训练吞吐量的技术。

---

十、内存优化：模型训练的主要内存组成

10.1 训练内存由哪些部分构成

深度学习训练的设备内存通常包含：

[
M_{\text{total}}

M_{\text{parameters}}
+
M_{\text{gradients}}
+
M_{\text{optimizer}}
+
M_{\text{activations}}
+
M_{\text{workspace}}
]

其中：

• (M_{\text{parameters}})：模型参数；
• (M_{\text{gradients}})：梯度；
• (M_{\text{optimizer}})：优化器内部状态；
• (M_{\text{activations}})：前向计算产生并需要保留的中间激活；
• (M_{\text{workspace}})：算子工作区、临时缓冲区及其他运行时内存。

实际设备内存还可能包含编译器分配、缓存与内存池等开销，因此这一表达式只是概念性的内存分解。

10.2 参数与优化器状态

假设模型包含：

[
P
]

个参数。

如果每个参数以 FP32 存储，那么模型权重理论上需要：

[
M_{\text{parameters}}=4P
]

字节。

如果同时保留 FP32 梯度：

[
M_{\text{gradients}}=4P
]

对于常见 Adam 实现，其一阶矩和二阶矩状态通常也采用 FP32，因此需要近似：

[
M_{\text{optimizer}}=8P
]

字节。

忽略其他开销时，三者合计为：

[
M\approx16P
]

字节。

例如，对于 1 亿参数的模型：

[
M\approx1.6\times10^9
]

字节，即约 1.49 GiB。

这还没有包含中间激活、计算工作区、数据缓冲区及其他运行时状态。

因此，不能仅根据模型参数文件的大小判断训练时的实际显存或 HBM 需求。

10.3 激活内存

在神经网络训练中，反向传播通常需要使用前向传播阶段的部分中间结果。

因此，某些激活张量需要保留到对应的梯度计算完成。

对于一个包含 (L) 个网络层、Batch Size 为 (B)、序列长度为 (T)、隐藏维度为 (H) 的模型，激活内存可以粗略理解为随以下规模增长：

[
M_{\text{activations}}
\propto
B\times T\times H\times L
]

这是用于理解规模变化的近似关系。

实际 Transformer 还可能包含与注意力矩阵相关的中间张量，其内存复杂度受到具体注意力实现方式影响。

因此，实际峰值内存必须结合模型计算图和算子实现分析。

10.4 混合精度为什么不能减少所有内存

混合精度通常降低部分激活数据和中间计算结果的存储成本。

但如果模型权重仍以 FP32 保存，那么参数存储本身就不会因为计算采用 BF16 而直接减半。

此外，Adam 等优化器的内部状态也可能继续使用 FP32。

因此，混合精度的内存收益取决于：

• 激活内存在总内存中的占比；
• 参数与优化器状态规模；
• 算子临时内存；
• 后端的内部缓存与计算策略。

对于激活内存占比较高的模型，混合精度可能产生明显收益。

对于主要由参数和优化器状态占用内存的模型，收益可能相对有限。

---

十一、Rematerialization：以计算换取内存

11.1 重计算的基本原理

重计算（Rematerialization），也常被称为 Activation Checkpointing 或 Gradient Checkpointing，是一种以额外计算换取内存空间的技术。

在普通反向传播中，系统可能需要保存大量前向计算的中间结果。

例如：

Forward:

X → Layer 1 → A1 → Layer 2 → A2 → Layer 3 → A3
                │             │             │
                └── 保存 ──────┴── 保存 ──────┘

在反向传播期间，这些中间结果可直接用于梯度计算。

重计算则采用不同策略。

它只保留部分中间状态，在反向传播需要其他激活值时，重新执行对应的前向计算。

Forward:

X → Layer 1 → Layer 2 → Layer 3
│                          │
└── 保留必要检查点 ──────────┘

Backward:

重新计算部分中间激活
            │
            ▼
         计算梯度

这种方式可以显著降低部分模型的峰值激活内存，但会增加额外的计算量。

11.2 时间与内存之间的权衡

假设普通训练需要：

[
M_{\text{normal}}
]

的峰值内存。

启用重计算后：

[
M_{\text{remat}}<M_{\text{normal}}
]

但训练计算量通常会增加：

[
F_{\text{remat}}>F_{\text{normal}}
]

因此，重计算并不是一种保证提高速度的优化技术。

它的主要作用是降低内存需求。

只有当降低内存占用允许使用更大的 Batch、更长的序列或避免昂贵的内存溢出处理时，重计算才可能间接提高整体训练效率。

11.3 Keras 的 remat()

Keras 3 提供：

keras.remat()

用于对指定函数应用重计算。

例如：

import keras
from keras import ops


def transform(x):
    for _ in range(3):
        x = ops.tanh(x + 0.1 * x)

    return x


remat_transform = keras.remat(transform)

随后：

y = remat_transform(x)

在支持的执行环境中，反向传播可以重新计算部分中间结果，而不是保留全部激活。

这个示例使用无状态数学操作，目的是展示重计算函数的基本结构。

对于包含随机操作、可变状态或其他副作用的复杂函数，还需要检查其在目标后端中的重计算语义。

11.4 RematScope

Keras 还提供：

keras.RematScope

用于在指定作用域中控制重计算策略。

例如：

inputs = keras.Input(shape=(128,))

with keras.RematScope(
    mode="list_of_layers",
    layer_names=["large_dense"]
):
    x = keras.layers.Dense(
        512,
        activation="gelu",
        name="large_dense"
    )(inputs)

outputs = keras.layers.Dense(
    1,
    dtype="float32"
)(x)

model = keras.Model(
    inputs,
    outputs
)

该示例使用 list_of_layers 模式，针对指定网络层应用重计算策略。

Keras 官方还提供其他模式：

模式	作用
full	对支持的操作广泛应用重计算
activations	对相关激活操作应用重计算
larger_than	根据输出大小阈值选择重计算
list_of_layers	指定需要重计算的层
None	禁用重计算

需要注意，RematScope 应在相关层的实际调用阶段生效，而不仅仅是在创建 Layer 对象时使用上下文管理器。

11.5 重计算的使用原则

并不是所有层都适合重计算。

例如，如果某个计算模块本身执行成本很高，但只保存少量中间激活，那么重计算可能带来较大的计算开销，却只节省很少内存。

相反，如果某个模块产生大量需要保留的激活数据，那么重计算可能具有更好的内存收益。

因此，应当优先分析激活内存占用较高的模块，再决定应用范围。

对于包含大量 Transformer Block 的模型，可以考虑以 Block 为单位设计重计算边界，而不是无条件对每一个简单操作单独应用重计算。

---

十二、梯度累积与重计算的区别

梯度累积和重计算都与设备内存有关，但二者的作用机制不同。

特征	梯度累积	重计算
核心目的	使用较小 Micro-batch 获得较大有效批大小	减少中间激活内存
主要方法	多次计算梯度后统一更新	反向传播时重新执行前向计算
是否改变更新频率	是	通常不改变
是否增加重复前向计算	不以此为主要机制	是
是否减少单次激活需求	通过较小 Micro-batch 实现	通过不保存部分激活实现
主要代价	更新频率与训练执行组织发生变化	增加计算量

例如，假设一个模型使用 Batch Size 为 256 时发生内存不足。

可以首先尝试将 Batch Size 降为 64，并设置：

gradient_accumulation_steps=4

使有效批大小保持在 256。

如果 Batch Size 为 64 时仍然因为激活内存不足而无法训练，则可以进一步考虑重计算。

这两种技术可以组合使用，但不应被视为相同的优化手段。

---

十三、性能分析：如何确定真正的瓶颈

13.1 为什么必须使用 Profiler

深度学习训练性能优化中，一个常见问题是，仅凭设备利用率或单步耗时判断瓶颈。

例如，GPU 利用率较低，可能由多种原因造成：

• 数据加载速度不足；
• 模型计算规模太小；
• Python 调度开销较高；
• 编译器生成了低效执行图；
• 模型频繁触发同步；
• 大量时间用于设备之间的数据传输。

不同原因需要完全不同的优化方法。

因此，在调整 Batch Size、精度或分布式配置之前，应优先分析实际执行时间分布。

13.2 Step Time 与吞吐量

假设模型在稳定训练阶段，平均每个 Step 耗时：

[
T_{\text{step}}=0.05s
]

每个 Step 处理 128 个样本。

那么样本吞吐量为：

[
P=\frac{128}{0.05}=2560
]

即：

[
2560\ \text{Samples/s}
]

需要注意，梯度累积会改变参数更新频率，因此应当区分：

• 每秒处理的样本数量；
• 每秒执行的训练批次数量；
• 每秒完成的优化器更新次数。

在比较训练配置时，这三个指标具有不同意义。

13.3 MFU

MFU（Model FLOPs Utilization）用于衡量模型计算量相对于硬件理论计算峰值的利用程度。

其简化表达式为：

[
\operatorname{MFU}

\frac{
F_{\text{model}}\times P_{\text{samples}}
}{
F_{\text{peak}}
}
]

其中：

• (F_{\text{model}})：处理一个样本所需要的模型浮点运算量；
• (P_{\text{samples}})：每秒处理的样本数量；
• (F_{\text{peak}})：相应硬件的理论浮点运算峰值。

对于多设备训练，分母应使用参与计算的设备所对应的理论总峰值，并保持计算精度口径一致。

MFU 可以帮助判断模型是否充分利用硬件计算能力。

但它并不等价于硬件监控界面中的 GPU Utilization、Tensor Core Utilization 或 TPU 活跃时间比例。

例如，一个模型可能在大部分时间内都有设备计算活动，但由于矩阵规模较小、计算访存比例较高或其他限制，实际达到的 FLOP/s 仍然较低。

因此，不同利用率指标不能直接相互替代。

13.4 JAX Profiler

对于 JAX 后端，可以使用 JAX Profiler 采集运行时执行信息。

例如：

import jax

jax.profiler.start_trace(
    "/tmp/jax_profile"
)

# 执行需要分析的训练计算
history = model.fit(
    x_train,
    y_train,
    epochs=2,
    verbose=0
)

jax.profiler.stop_trace()

生成的 Trace 可以使用 XProf 等工具进一步分析。

但在实际性能分析中，应尽量将首次编译阶段与稳定训练阶段分开采集。

否则，编译耗时可能掩盖真实的设备计算性能。

此外，JAX 通常采用异步设备执行机制。

对于直接调用 JAX 计算函数的基准测试，需要通过适当的设备同步确保测量包含实际计算完成时间。

否则，测量结果可能主要反映计算任务提交的耗时，而不是实际执行耗时。

13.5 TensorFlow 与 PyTorch 的性能分析

对于 TensorFlow 后端，可以使用 TensorFlow Profiler 分析：

• TensorFlow 运算时间；
• GPU 内核执行；
• 输入管道性能；
• 设备调度；
• 内存使用。

对于 PyTorch 后端，可以使用 torch.profiler 分析：

• CPU 与 CUDA 执行时间；
• 算子调用；
• 计算图行为；
• 内存分配；
• 内核执行。

虽然 Keras 提供统一的模型接口，但底层性能分析通常仍需要使用后端相应的 Profiler。

这是 Keras 多后端架构的一个重要边界。

---

十四、性能优化应当遵循的顺序

在实际工程中，不应将所有性能优化技术同时启用，然后仅根据最终速度判断它们是否有效。

这种方式会使各项优化之间的因果关系难以分析。

更合理的方法是建立基线，并逐步调整。

14.1 第一步：建立基线

首先使用简单、稳定的配置训练模型。

例如：

keras.config.set_dtype_policy(
    "float32"
)

model.compile(
    optimizer="adam",
    loss="mse",
    jit_compile="auto"
)

记录：

• 模型参数数量；
• Batch Size；
• 稳定训练吞吐量；
• 峰值内存；
• 编译耗时；
• 训练指标。

基线不仅用于性能比较，也用于验证后续优化是否改变模型行为。

14.2 第二步：分析瓶颈类型

可以将瓶颈大致划分为以下类型。

瓶颈类型	主要表现	优先分析方向
数据供应不足	设备频繁等待输入	数据加载、预取、缓存
Python 调度开销高	单步计算量小、调度频繁	JIT、执行步骤组合
计算能力利用不足	矩阵计算吞吐量较低	Batch、张量形状、算子
内存带宽限制	数据访问占据主要耗时	算子融合、布局、精度
内存容量不足	OOM、Batch 无法增大	混合精度、累积、重计算
多设备通信限制	增加设备后加速不足	并行策略、通信布局

这张表只提供初步排查方向，不能替代实际的性能采样与分析。

14.3 第三步：调整数值精度

对于支持 BF16 或 FP16 加速的设备，可以优先尝试混合精度。

例如，在 TPU 上：

keras.config.set_dtype_policy(
    "mixed_bfloat16"
)

然后重新测试吞吐量与训练数值表现。

如果模型出现异常，需要检查具体运算和输出精度，而不是简单认为所有低精度训练都不可靠。

14.4 第四步：调整 Batch Size

在设备内存允许的范围内，可以测试不同 Batch Size。

例如：

32
64
128
256
512

观察每种配置对应的训练吞吐量和内存使用。

但不能只选择每秒处理样本最多的配置。

如果训练算法对 Batch Size 敏感，还需要结合验证指标和收敛速度进行判断。

14.5 第五步：优化编译与调度

如果模型受运行时调度限制，可以尝试：

jit_compile=True

并在支持的后端中测试：

steps_per_execution=8

或者其他执行步数。

对于 JAX 和 TPU，还应检查输入形状是否导致重复编译。

14.6 第六步：调整内存与分布式策略

如果单设备内存不足，可以依次考虑：

• 减少 Micro-batch；
• 使用梯度累积；
• 对高内存模块启用重计算；
• 在必要时采用模型并行。

如果模型能够正常放入单设备，但训练时间过长，则可以评估数据并行。

需要注意，这些策略并不存在对所有模型都适用的固定顺序。

实际决策应根据 Profiler 和设备内存数据确定。

---

十五、完整实例：在 JAX 后端中使用 Keras 高性能训练

下面通过一个完整示例，将混合精度、JIT 编译、梯度累积及多设备数据并行组合起来。

该示例以 JAX 后端为基础，适用于已经正确安装 JAX 和 Keras 的环境。

在 TPU 上运行时，可以启用 BF16；如果在 CPU 环境中测试，则建议改用 FP32。

15.1 环境配置

import os

# 必须在导入 Keras 前设置
os.environ["KERAS_BACKEND"] = "jax"

import numpy as np
import keras
import jax

from keras import layers

print("Backend:", keras.config.backend())
print("Devices:", jax.devices())

对于 TPU 环境：

keras.config.set_dtype_policy(
    "mixed_bfloat16"
)

如果需要使用编译缓存：

jax.config.update(
    "jax_compilation_cache_dir",
    "/tmp/jax_compile_cache"
)

实际训练前，应确认硬件运行时支持对应的数据类型和 JAX 版本。

15.2 配置数据并行

如果环境具有多个可用加速器，可以设置：

devices = keras.distribution.list_devices()

if len(devices) > 1:
    distribution = keras.distribution.DataParallel(
        devices=devices
    )

    keras.distribution.set_distribution(
        distribution
    )

这里需要在创建模型变量之前完成分布式配置。

如果只有一个设备，则保留默认的单设备执行方式。

15.3 构建训练数据

rng = np.random.default_rng(42)

x = rng.normal(
    size=(4096, 64)
).astype("float32")

true_w = rng.normal(
    size=(64, 1)
).astype("float32")

noise = rng.normal(
    scale=0.1,
    size=(4096, 1)
).astype("float32")

y = x @ true_w + noise

这里构建了 4,096 个样本，每个样本具有 64 个输入特征。

为了使示例不依赖外部文件，目标使用线性函数与随机噪声生成。

15.4 构建模型

inputs = keras.Input(
    shape=(64,)
)

x_model = layers.Dense(
    256,
    activation="gelu"
)(inputs)

x_model = layers.Dense(
    256,
    activation="gelu"
)(x_model)

x_model = layers.Dense(
    128,
    activation="gelu"
)(x_model)

outputs = layers.Dense(
    1,
    dtype="float32"
)(x_model)

model = keras.Model(
    inputs,
    outputs
)

模型主体采用全局 BF16 混合精度策略。

最后一层显式采用 FP32，以便输出回归预测值时保留较高精度。

15.5 配置训练系统

optimizer = keras.optimizers.AdamW(
    learning_rate=1e-3,
    weight_decay=1e-4,
    gradient_accumulation_steps=4
)

model.compile(
    optimizer=optimizer,
    loss=keras.losses.MeanSquaredError(),
    metrics=[
        keras.metrics.MeanAbsoluteError()
    ],
    jit_compile="auto",
    steps_per_execution=8
)

这套配置包含四项性能相关机制。

第一，模型主体采用混合 BF16 精度。

第二，JAX 在支持的情况下使用 XLA 编译。

第三，多个 Batch 可以组合为一次编译函数执行。

第四，优化器在累计 4 个训练 Batch 的梯度后执行参数更新。

需要强调，梯度累积次数为 4，并不意味着 steps_per_execution=8 会把 8 个 Batch 的梯度合并成一次更新。

两者控制的是不同的执行维度。

15.6 执行训练

history = model.fit(
    x,
    y,
    batch_size=128,
    epochs=10,
    validation_split=0.1,
    shuffle=True,
    verbose=2
)

在该示例中，训练数据由 NumPy 提供。

如果使用多设备 DataParallel，Keras 将按照分布式配置分配全局 Batch。

这里的 batch_size=128 表示传入训练系统的全局批大小，而不是每个设备分别处理 128 个样本。

由于设置了：

gradient_accumulation_steps=4

在批次大小与梯度归约条件满足要求的情况下，每次参数更新对应的有效样本数量可以近似理解为：

[
128\times4=512
]

但实际参数更新次数还取决于每个 Epoch 的批次数量以及不完整批次的处理方式。

15.7 检查精度策略

可以查看网络层的精度配置：

for layer in model.layers:
    print(
        layer.name,
        layer.compute_dtype,
        layer.variable_dtype
    )

对于采用全局混合精度策略的 Dense 层，通常可以看到 BF16 计算类型和 FP32 变量类型。

而最后的输出层采用 FP32。

15.8 性能分析

正式进行性能评估时，应避免将首次编译时间直接纳入稳定吞吐量比较。

例如，可以分别记录：

• 第一次 fit() 的总耗时；
• 后续相同配置的稳定训练耗时；
• 不同 Batch Size 下的 Samples/s；
• 是否发生重复编译；
• 峰值设备内存；
• 验证指标变化。

对于多设备训练，还应进一步检查设备间通信时间和分布式扩展效率。

需要强调，本示例的重点是展示 Keras 3 高性能计算接口之间的组合关系。

它并不保证启用所有这些功能后，训练速度一定比基础配置更快。

实际收益仍然需要通过目标硬件和具体模型的性能测试验证。

---

十六、高性能训练中的常见设计问题

16.1 认为混合精度一定能够提高训练速度

混合精度能否加速，取决于硬件能力和具体算子。

如果模型主要包含内存访问、数据转换或 Python 控制逻辑，那么低精度矩阵运算带来的收益可能有限。

此外，低精度与高精度之间的频繁转换也可能增加额外开销。

因此，混合精度应通过实际吞吐量和数值稳定性测试验证。

16.2 认为 JIT 编译必然降低总训练时间

JIT 编译需要支付初始编译成本。

对于短时间任务，编译成本可能超过后续执行收益。

对于输入形状频繁变化的模型，还可能产生多次编译。

因此，应区分首次执行性能与稳定训练性能。

16.3 认为 GPU 利用率高就意味着训练效率高

设备利用率、Tensor Core 活跃度和 MFU 衡量的是不同对象。

即使设备持续执行计算，模型也可能因为算子规模、访存或执行布局等原因，无法达到较高的计算吞吐量。

因此，不能仅根据单一设备利用率指标判断训练系统已经达到最优状态。

16.4 认为梯度累积能够直接加速训练

梯度累积的主要作用，是在较小 Micro-batch 下获得较大的有效批大小。

它通常不会凭空增加硬件计算能力。

相反，如果单次 Micro-batch 过小，还可能降低矩阵计算效率，并增加训练调度成本。

因此，梯度累积应主要用于满足内存或优化算法需求。

16.5 认为增加设备数量能够获得线性加速

分布式训练存在通信成本、同步成本和负载不均衡问题。

当模型过小或每个设备上的计算量不足时，通信成本可能占据较大比例。

因此，增加设备数量后，应通过实际并行效率判断扩展是否合理。

16.6 认为重计算是一种通用加速方法

重计算以额外计算换取较低激活内存。

如果模型本身已经受到计算能力限制，而内存并不紧张，那么启用重计算通常可能降低训练速度。

因此，只有在内存压力确实影响模型规模或批次选择时，重计算才更有意义。

16.7 在没有性能基线的情况下同时修改多个参数

如果同时改变 Batch Size、精度、JIT、梯度累积和分布式配置，那么即使训练速度发生变化，也难以确定真正起作用的因素。

因此，应尽可能采用控制变量的实验方法。

每次调整都应记录配置、运行环境、吞吐量、内存和训练指标，使性能分析具有可重复性。

---

十七、Keras 高性能计算系统的架构意义

从整个 Keras 3 框架来看，高性能计算并不是位于某个独立位置的单一模块。

它实际上横跨模型、训练、数据处理和计算后端等多个层次。

可以将其概括为：

                 Keras Model
                      │
                      ▼
              Precision Policy
                FP32 / FP16
                   BF16
                      │
                      ▼
                Training Loop
                      │
        ┌─────────────┼─────────────┐
        │             │             │
        ▼             ▼             ▼
      JIT         Batch Control   Remat
   Compilation     Accumulation   Memory
        │             │             │
        └─────────────┼─────────────┘
                      ▼
                Distribution
               Data / Model
                      │
                      ▼
                Backend Runtime
          ┌───────────┼───────────┐
          ▼           ▼           ▼
         JAX      TensorFlow    PyTorch
          │           │           │
          └───────────┼───────────┘
                      ▼
              Hardware Execution
               CPU / GPU / TPU

这一结构体现了 Keras 的架构定位。

Keras 负责提供统一的模型定义与部分性能控制接口。

计算后端负责将计算图和张量操作转换为实际执行程序。

硬件运行时则负责在具体设备上执行计算。

从软件工程角度看，这种设计具有三个重要意义。

第一，模型数学结构与性能配置可以相对分离。

在不改变神经网络基本结构的情况下，开发者可以通过精度策略、编译配置和分布式策略调整执行方式。

第二，性能优化可以在不同层次独立进行。

例如，数据管道负责提高输入供应效率，JIT 编译负责降低计算调度成本，混合精度负责优化部分数值计算，分布式策略则负责扩展硬件资源。

第三，性能配置仍然受到后端能力的约束。

Keras 的统一接口并不意味着所有后端拥有相同的编译器、相同的设备通信机制或相同的内存管理策略。

因此，在深入优化时，仍然需要理解具体后端的执行机制。

换言之，Keras 主要降低了高性能计算能力的使用门槛，但并没有消除底层计算系统本身的复杂性。

结语

Keras 3 的高性能计算系统，主要解决如何在保证训练目标和数值正确性的前提下，更高效地使用 CPU、GPU 和 TPU 等计算资源。

其核心并不是某一项单独的加速技术，而是通过数值精度、计算编译、分布式执行、批次组织和内存优化之间的协同，提高整个训练系统的效率。

混合精度通过降低部分计算与存储的数值位宽，利用现代加速器的低精度计算能力。

JIT 编译通过优化计算图和减少运行时调度开销，提高计算执行效率。

数据并行通过复制模型，将不同数据分配到多个设备执行。

模型并行通过对参数和中间张量进行分片，支持更大规模的模型训练。

梯度累积允许使用较小的 Micro-batch 获得较大的有效批大小。

重计算则通过重复执行部分前向运算，降低反向传播所需要的激活内存。

这些技术分别解决不同问题，因此不能简单地将它们全部视为通用加速方法。

有些技术主要提高计算吞吐量，有些技术主要降低内存需求，还有些技术主要改变参数更新的组织方式。

对于实际模型训练，最重要的是首先明确瓶颈所在，再选择与瓶颈对应的优化手段。

如果模型主要受到数据管道限制，就应当优先优化数据加载。

如果主要受到 Python 调度限制，可以考虑 JIT 编译和执行步骤组合。

如果内存不足，可以考虑梯度累积与重计算。

如果单设备训练时间过长，且计算规模足以覆盖通信成本，则可以进一步评估分布式训练。

因此，高性能训练不应被理解为一套固定参数配置，而应被理解为一个持续测量、分析与优化的工程过程。

至此，我们已经介绍了 Keras 3 的基础计算、神经网络层、模型构建、训练、数据处理、训练控制和高性能计算七个主要部分。

这些组件共同形成了一个能够定义、训练、扩展和优化深度学习模型的完整计算体系。

但当模型训练完成后，还需要解决另一个重要问题：如何保存训练结果、恢复模型状态，并将模型转换为适合实际应用环境的推理形式。

在下一篇文章中，我们将介绍 Keras 3 的模型持久化与部署系统，重点分析模型序列化、权重保存、检查点恢复、模型导出与推理部署之间的关系，并讨论不同计算后端下的模型兼容性与部署边界。