深入探究适用于大型嵌入模型 (LEM) 的 SparseCore

SparseCore 是一款专门的平铺处理器，旨在为涉及不规则、稀疏内存访问和计算的工作负载提供高性能加速，尤其是在存储在高带宽内存 (HBM) 中的大型数据集上。虽然它在嵌入查找等任务中表现出色，但其功能还可扩展到加速各种其他动态和稀疏工作负载。

1. SparseCore 简介

主要架构特性：

  • 平铺式架构：包含多个计算块（每个块都是一个完整的数据流单元，具有自己的本地内存和处理单元），可实现并行处理。
  • 动态执行：原生支持依赖于数据控制流和内存访问，对于稀疏数据至关重要。
  • 矢量处理：利用小矢量任务（8 元素或 16 元素，具体取决于硬件版本）进行高效计算。
  • 集中式控制：单个 SparseCore 序列器可编排所有功能块的任务，确保操作同步。
  • 数据汇总支持：包括专门的跨通道操作，有助于执行排序、过滤和前缀和等任务。
  • 内存层次结构：战略性地利用 HBM 来存储大型数据集，并利用本地暂存内存 (SPMEM) 来暂存频繁访问的数据，从而显著缩短 HBM 延迟时间。

规格一览：

属性 TPU v4 TPU v5p Trillium
SparseCore/芯片 4 4 2
Tiles/SparseCore 16 16 16
SIMD 宽度 8 8 8 (F32) 16 (BF16)
HBM 容量 32 GiB 96 GiB 32 GiB

2. SparseCore 主机预处理

有效的数据准备对于 SparseCore 性能至关重要，而这正是主机预处理发挥重要作用的地方。其中包含几项关键功能：

  • 数据转换：
    • 对原始输入数据应用必要的转换。
    • 管理 ID 转换，这在处理特征或表格堆叠时尤为重要。
    • 将输入数据转换为坐标 (COO) 稀疏格式，详见下文。
    • 对数据进行分区，以便在芯片上可用的不同 SparseCore 之间高效分配数据。
  • 限制验证：
    • 确保输入数据的特征（例如 ID 数量）符合 SparseCore 的预定义操作限制，例如 max_ids_per_partition 和 max_unique_ids_per_partition。
    • 如果输入数据超出这些限制，宿主预处理层可以尝试将数据分段为符合限制的较小微批次。
  • 数据传输：
    • 高效地将处理和验证后的数据复制到 TPU 的高带宽内存 (HBM)，使其可用于 SparseCore 执行。

了解表格堆叠：

表堆叠是一种重要的优化技术，可将多个嵌入表在逻辑上组合在一起，以提高嵌入查找效率。此过程通常由底层机器学习框架自动处理。

  • 特征堆叠：当多个不同的特征共享同一个底层嵌入表时，就会发生这种情况。一个常见示例是使用单个嵌入字典来处理各种分类特征，例如来自不同上下文的邮政编码。
  • 表堆叠：在此方案中，多个不同的嵌入表堆叠在一起。共享相同嵌入维度和优化器配置的表通常会分组。

堆叠表的主要优势在于，可以为对这些堆叠表执行的操作创建更大的有效批次大小。这可以减少计算开销，并有效地隐藏芯片间通信 (ICI) 延迟。为获得最佳性能，建议堆叠适量的表（一般在 5 到 100 个之间）。

3. 转换为 COO 张量

在 SparseCore 处理数据之前，通常会将其转换为坐标 (COO) 稀疏张量格式。COO 格式是一种高效表示稀疏矩阵的方式，通常使用三个数组：

  • row_ids：一个数组，包含每个非零元素的行索引。在批处理的上下文中，这通常对应于批次维度。
  • col_ids：一个数组，包含每个非零元素的列索引。 对于嵌入，这些通常是特征或 ID 值。
  • values（可选）：一个数组，用于保存相应 (row, col) 坐标处非零元素的实际值。对于与 ID 数量相关的限制计算（稍后讨论），通常不考虑这些值（增益）。

说明性示例：

假设有一个表示 ID 批次的输入稀疏矩阵：

[
    [id_A],                 // Sample 0
    [id_A, id_B, id_C],     // Sample 1
    [id_B, id_B, id_D],     // Sample 2 (note duplicate id_B)
]

转换为 COO 格式后（可能还会对同一样本中的 ID 进行去重）：

row_ids = [0, 1, 1, 1, 2, 2]
col_ids = [id_A, id_A, id_B, id_C, id_B, id_D]

此转换对于 SparseCore 处理和分配工作至关重要。尤其是 col_ids，对于确定 ID 所属的特定 SparseCore 分区至关重要，可实现高效的分片和查找。

4. SparsecoreConfig：高级别 API

特定于框架的嵌入 API：

SparsecoreConfig 或 XLA 标志等同等机制可作为高级接口来控制各种 SparseCore 行为。深入了解这些参数对于有效调整性能和确保模型正常运行至关重要。

  • disable_table_stacking: bool = False
    • 说明：此标志用于控制是否自动堆叠表格，以防止框架堆叠表格，这可能会因开销增加和隐藏芯片间互连 (ICI) 延迟的能力下降而导致性能降低。
    • 默认：False（表示在框架支持的情况下，默认情况下通常会启用表格堆叠）。
  • max_ids_per_chip_per_sample: int = 64
    • 说明：此形参用于设置单个芯片可以从输入批次中的一个样本处理的嵌入 ID 总数的全局上限（汇总了所有表）。它是一种在考虑更精细的单表或单分区限制之前，在芯片级管理资源的机制。对该值进行微调通常取决于具体的模型特征和整体系统容量。
    • 默认值：64。
  • max_ids_per_table: Optional[Dict[str, int]] = None
    • 说明：此形参用于指定每个逻辑表可处理的嵌入 ID（可包含重复项）数量上限，其中考虑了所有 SparseCore 中该逻辑表的所有分区。此限制比 max_ids_per_partition 更宽松。如果表 T 分为 P 个分区，则此限制适用于定向到所有 P 分区的 ID 的总和。它通常与 max_ids_per_partition_per_sample 和整体批次大小有关。
    • 设置：通常使用限制文件（例如，使用 xla_sparse_core_max_ids_file 标志）进行配置，其中定义了 max_ids_per_partition。这种表级概念是一种设置分区级限制（max_ids 和 max_uniques）的方法。
    • 默认值：None（如果未明确提供，则可以从每个分区的限制或其他配置中推断出该值）。
  • max_unique_ids_per_table: Optional[Dict[str, int]] = None
    • 说明：与 max_ids_per_table 类似，但此参数指定了每个逻辑表的最大唯一 ID 数量。这是一项关键设置，用于适当调整在唯一 ID 处理和后续矢量操作中使用的设备端缓冲区的大小。
    • 设置：通常也在限制文件中定义，或从 max_unique_ids_per_partition_per_sample 派生。
    • 默认值：None。
  • allow_id_dropping: bool = False
    • 说明：此布尔值标志用于控制在输入数据中遇到的 ID 数量（观测到的限制）超过编译期间设置的限制（例如 max_ids_per_partition）时是否丢弃 ID。
      • 如果 True：会导致超出限制的 ID 会被静默舍弃。通常，系统会按排序顺序处理分区中的 ID，并舍弃任何会使指定小批次的运行计数超过限制的 ID。这样一来，程序便可继续执行，但可能会对模型准确性产生不利影响。
      • 如果为 False：如果观测到的限制超出编译的限制，系统会触发错误，并且进程可能会终止。此方法可确保处理所有数据，但需要更保守地配置限制。
    • 默认值：False（导致溢出时出错，而不是静默丢弃数据）。

  • initialize_tables_on_host: bool = True

    • 说明：此标志用于确定是否在主机 CPU 上初始化嵌入表，然后再将其转移到 TPU 的高带宽内存 (HBM)。标准做法是在主机上初始化表。将此值设置为 True 符合此惯例。 如果设置为 False，则表示采用设备端初始化机制，这可能会对性能产生不同的影响，或者有特定的初始化前提条件。

  • enable_fast_table_initialization: bool = False

    • 说明：直接在 TPU 上初始化表。这有助于缩短模型启动时间。

5. 流水线处理以提升性能

流水线是一种性能优化技术，可实现 TensorCore (TC) 和 SparseCore (SC) 上操作的同步执行。通过重叠这些计算，可以显著提高整体吞吐量。

  • 机制：在涉及稀疏嵌入查找（由 SC 处理）和密集层计算（由 TC 处理）的标准训练步骤中，流水线处理允许 SC 处理步骤 i 的一部分（例如，前向或后向传递），而 TC 同时处理同一步骤 i 的另一部分，甚至相邻步骤（如 i-1 或 i+1）的部分。
  • 对梯度的影响：SparseCore 可能会对“过时”的梯度进行操作。 例如，在步骤 i 的反向传播阶段计算出的梯度可能不会完全更新，并且在步骤 i+2 之前不会对 SC 可见。
  • 性能与数值之间的权衡：这种重叠执行可以大幅提高速度，设备步进时间最多可缩短 2 倍。不过，使用过时梯度导致的数值（嵌入权重）细微变化可能会影响模型收敛行为或最终实现的准确性。这种权衡的合理性高度依赖于模型，通常需要进行实证验证。
  • 控制标志：流水线可以通过 tf_xla_disable_full_embedding_pipelining 进行控制。将此标志设置为 true 会停用完全流水线处理（重叠的 TensorCore 和 SparseCore 计算），而将其设置为 false（或如果标志的语义表示在为 false 时启用）则会激活它。

概念性流水线流程：

  • 不使用流水线（简化的顺序流程）：

    Loop: SC/F_i -> TC/F_i -> TC/B_i -> SC/B_i

  • 使用流水线（简化的重叠流程）：

    Time ->
Step i:   SC/F_i | TC/F_i | TC/B_i | SC/B_i
Step i+1:          SC/F_i+1| TC/F_i+1| TC/B_i+1| SC/B_i+1

    注意：在硬件和编译器中实现的实际流水线阶段可能更加复杂，通常涉及预循环、主执行循环和后循环，以管理数据依赖关系并确保正确性。

6. XLA 的作用

XLA（加速线性代数）是一种针对特定领域的编译器，可将高级计算图（通常来自 TensorFlow 等框架）转换为专为 TPU 量身定制的高度优化的机器码。这包括为 SparseCore 生成操作指令。

SparseCore 上下文中的关键函数：

  • 稀疏操作的编译：XLA 负责将嵌入查找操作（例如 SparseDenseMatmulOp）和其他稀疏计算编译为低级可执行的 SparseCore 程序。
  • 限制集成：它利用配置的操作限制（例如 max_ids_per_partition、max_unique_ids_per_partition，通常通过标志（如 xla_sparse_core_max_ids_file）指定的文件提供）来静态确定设备上内存缓冲区的大小并分配这些缓冲区，尤其是在 SPMEM 中。
  • 有针对性的优化：XLA 会执行一系列专门针对 SparseCore 架构设计的优化。这些优化包括指令调度、内存布局转换和操作融合，以最大限度提高效率。
  • 使用标志进行控制：SparseCore 行为、调优参数和优化策略的许多方面都通过 XLA 标志（例如用于限制估计的 xla_sparse_core_estimate_max_ids 或用于调试的 xla_sc_detect_nan）公开和控制。

开源状态：

目前，Sparsecore 实现是内部实现，通过 libtpu.so 提供。

错误报告和诊断：

与 SparseCore 配置或资源限制相关的编译失败通常表现为 XLA:TPU 编译时错误。这些错误消息可以提供有关问题的宝贵数据洞见，例如为可用 SPMEM 设置的限制过高，或使用了不受支持的配置。

7. 限制如何转换为 SparseCore 上的表格

在 SparseCore 上，“限制”是基本配置参数，主要指每个分片（分布）在可用 SparseCore 上的表的两个分区级设置：

  • max_ids_per_partition：此参数用于定义在单个计算步骤中，任何单个 SparseCore 预计会发送到或处理给定表的特定分区中的 ID 总数（包括重复项）上限。
  • max_unique_ids_per_partition：此属性定义了任何单个 SparseCore 预期会发送到或处理的唯一 ID 的数量上限。

转换为物理表布局和处理：

  • 表分片策略：嵌入表通常在系统中的所有 SparseCore 之间进行“mod-sharded”。这意味着，每个 SparseCore 都负责每个表的不同词汇子集（行）。系统通常会根据 k = j % num_total_sparse_cores 等公式为 SparseCore_k 分配 ID j。
  • “分区”的定义：在此上下文中，“分区”是指嵌入表中的特定区段，单个 SparseCore 可处理该区段的查找。
  • SPMEM 缓冲区分配：XLA 编译器使用这些限制在设备端暂存区内存 (SPMEM) 中静态调整缓冲区大小并分配缓冲区。缓冲区的尺寸经过调整，可将给定分区 ID 的所有必要数据（最多达到指定的 max_ids 和 max_unique_ids 限制）加载到 SPMEM 中进行处理。对于非逐元素计算（例如，在创建压缩稀疏行 [CSR] 表示形式时，减少分区内的重复 ID），这一点尤为重要，因为该分区 ID 的整个相关数据集需要随时在快速内存中可用。

  • 编译限制与观测限制：

    • 观测到的限制：这些是在运行时根据正在处理的输入数据为每个分区遇到的实际 ID 数量。
    • 如果观测到的限制超出编译的限制，则可能会导致 ID 丢弃（如果 allow_id_dropping 已启用）或错误。

  • 计算限值：确定适当限值的过程需要仔细分析输入数据分布。对于任何给定的表格（我们将其称为 T1，它本身可能是更大的堆叠表格 T 的一部分）：

    1. 输入批次（例如，形状为 [BatchSize, MaxSequenceLength] 的二维 SparseTensor）最初在可用的 SparseCore 之间进行拆分。例如，如果 TensorCore 与 2 个 SparseCore 配对，则每个 SparseCore 可能会收到形状为 [BatchSize/2, MaxSequenceLength] 的子批次。
    2. 然后，将此子批次转换为 COO 格式，从而生成 row_ids 和 col_ids。
    3. 系统会移除同一样本中的重复 ID（即具有相同 row_id 和 col_id 的条目）。
    4. 对于每个剩余的唯一 col_id（在样本内），负责此 ID 的目标 SparseCore 使用 mod-sharding 规则确定：target_sc_id = col_id % num_total_sparse_cores。
    5. 系统会维护每个 target_sc_id 的 ID 总数 (ids_per_sparse_core[target_sc_id]++) 和唯一 ID 数 (unique_ids_per_sparse_core[target_sc_id]++，在确保特定 target_sc_id 的唯一性之后) 的计数。
    6. 然后，将表 T1 的 max_ids_per_partition 设置为 max(ids_per_sparse_core_array)。
    7. 同样，表 T1 的 max_unique_ids_per_partition 设置为 max(unique_ids_per_sparse_core_array)。
    8. 如果表 T1 是堆叠表的组成部分，则在汇总所有组成表的统计信息之前，可能会对 ID 分布应用其他转换（例如旋转或平移）。这有助于平衡芯片之间的负载。

正确设置这些限制是一项平衡措施：较低的限制可能会带来更高的性能（因为每步需要处理的数据更少，并且 SPMEM 压力会降低），但如果设置得过低，则可能会导致过多的微批处理或不必要的 ID 丢弃。

8. 每个 SparseCore 的通信方式

SparseCore 通信（尤其是在处理 ID 列表以进行嵌入查找时）依赖于多种协调机制：

  • 模块分片和隐式路由：
    • 嵌入表在系统中的所有 SparseCore 之间进行 mod-sharding。
    • 当宿主提供一批输入数据（随后预处理为 COO 格式，包括 col_ids）时，col_id 值用于确定哪个 SparseCore 负责该特定 ID：target_sc_id = col_id % num_total_sparse_cores。
    • 每个 SparseCore 实际上只会接收和处理映射到其分配的词汇分区中的 ID 子集。主机预处理阶段对于准备数据至关重要，这样每个 SparseCore 都可以轻松识别并处理其相关 ID。
  • 按主机划分的数据分布：
    • 主机预处理逻辑会对整体输入批次进行分区，并将 row_ids 和 col_ids 的相关部分（以及任何关联的特征或权重，如果适用）分配给每个 SparseCore 可直接访问的内存 (HBM)，或分配给 SparseCore 将从中提取所需数据的共享 HBM。
  • SparseCore 内部处理：
    • SparseCore 在收到指定的一组给定表分区的 ID 后，会执行这些 ID 的去重和收集相应嵌入向量等操作。这些主要是 SparseCore 自身 tile 内执行的本地计算，并利用其本地 SPMEM。
  • SparseCore 间通信（全部到全部）：
    • 在初始处理阶段（例如嵌入查找）之后，可能会使用“全到全”通信模式来合并或重新分配 SparseCore 之间的结果（例如，在将激活馈送到需要与所有原始样本位置对应的输入的 TensorCore 层之前）。如果原始输入批次是为并行处理而分布的，那么这对于重建完整的激活集至关重要。
  • 与 TensorCore 通信：
    • SparseCore 与 TensorCore 通信，以发送嵌入激活（在前向传播期间）并接收梯度（在反向传播期间）。此交互由 XLA 编译的程序精心安排，通常涉及 HBM 作为中间缓冲区。流水线策略（如前所述）对这种 SC-TC 通信的时序和同步有很大影响。

从本质上讲，将 ID 初始“分配”给相应的 SparseCore 主要由分片方案和宿主预处理步骤来处理。后续通信涉及 SparseCore 对其本地数据进行操作，如果需要在 TensorCore 进一步处理之前在 SparseCore 之间全局交换或重新排序数据，则可能需要进行全到全等集体通信操作。

9. SparseCore 内存管理

每个 SparseCore 都能高效管理多种不同类型的内存，以执行计算：

  • 暂存存储器 (SPMEM)：
    • Nature：一种相对较小但速度非常快的本地 SRAM，仅供每个 SparseCore 使用。请务必注意，SPMEM 不是缓存；它的使用由 XLA 编译器明确管理和编排。
    • 用途：SPMEM 用于“机会性地暂存数据”。这包括持续进行 SC 计算所需的输入、输出和中间结果。在 SPMEM 中暂存数据可显著减少通常与访问 HBM 相关的高延迟。
    • 大小调整：如“限制”部分中所述，SPMEM 缓冲区在编译时静态确定大小。此大小调整基于 max_ids_per_partition 和 max_unique_ids_per_partition 等参数。这种静态分配可确保对于表分区上的任何给定操作（例如 CSR 缩减），该分区 ID 的所有必要数据（不超过定义的限制）都可以放入 SPMEM 中。
    • 编译器优化：XLA 编译器采用复杂的优化技术，可精确确定需要在 SPMEM 中暂存多少数据以及哪些特定数据元素，从而有效隐藏 HBM 延迟并最大限度地提高性能。
    • 动态分配限制：SparseCore 编译器目前不支持动态暂存区分配。这突显了通过仔细配置限制来实现静态调整大小的关键重要性。
  • 高带宽内存 (HBM)：
    • 性质：一种大型共享内存资源，可供所有 SparseCore、TensorCore 和主机系统访问。主要嵌入表存储在 HBM 中。
    • 堆栈使用情况：SparseCore 操作通常需要在 HBM 中临时存储中间结果，这些结果要么无法放入有限的 SPMEM 中，要么需要在处理流水线的较大阶段之间传递。正向和反向传递期间的 HBM 堆栈使用情况可按如下方式估算：
      • 前向传递 HBM 栈（单表）≈ (2 * feature_width + 1) * max_unique_nz_per_row * logical_replica_count * 4 字节
      • 反向传递 HBM 栈（单表）≈ 3 * feature_width * max_unique_nz_per_row * logical_replica_count * 4 字节
    • 堆使用情况：HBM 还可容纳由主机管理的堆。堆存储数据，例如密集层权重、模型使用的常量和预提取的输入数据。堆使用量往往会随着主机预提取数据的步数（由 maximum_parallel_iterations 标志控制）而增加。虽然更多的预提取可以通过将主机到设备的传输与设备计算重叠来提高性能，但也会消耗更多 HBM。
    • HBM 优化序列化：标志 xla_sc_num_serialized_tables_to_optimize_hbm 提供了一种机制，用于控制在任意给定时间有多少个表的数据保持在 HBM 堆栈内存中。增加此数字可有效地串行化更多表的处理，从而减少 HBM 堆栈峰值使用量，但可能会因并行性降低而导致性能下降。
  • 矢量内存 (VMEM)：
    • VMEM 是 TC（TensorCore）专用的本地暂存内存。虽然 VMEM 不由 SparseCore 直接管理，但它是 SC 交互的内存生态系统的一个组成部分，主要通过 TensorCore 进行交互。

总体内存管理策略：

SparseCore 的核心内存管理策略是使用小型快速 SPMEM 来存储 SparseCore tile 正在积极处理的“热门”数据，从而最大限度地减少对较慢的 HBM 的访问。 配置的限制是确保 SPMEM 不会溢出的主要机制。HBM 用于存储大型嵌入表和临时数据，这些数据要么超出 SPMEM 容量，要么需要在不同的处理单元或流水线阶段之间共享。 XLA 编译器负责根据这些架构原则和用户配置的限制来协调所有数据移动和缓冲区分配。

10. 性能和内存瓶颈

若要通过 SparseCore 实现最佳性能，您需要清楚了解潜在的瓶颈以及如何解决这些瓶颈。这些问题可能出现在主机上、SparseCore 本身中，或出现在 SparseCore 与 TensorCore 的互动中。

常见性能瓶颈：

  • 主机瓶颈：
    • 问题：主机 CPU 可能无法快速预处理数据并将其馈送到 TPU，从而导致 SparseCore 和 TensorCore 利用率不足。这是一种常见的性能限制因素。
    • 缓解措施：监控主机 CPU 利用率和输入流水线指标。优化主机端数据加载和预处理例程（请参阅 COO 转化提示）。调整 maximum_parallel_iterations 标志以微调数据预提取。
  • 次优的 TC/SC 同步（缺少流水线）：
    • 问题：如果 TensorCore 和 SparseCore 之间的流水线被停用或无法高效运行，一个单元可能会花费大量时间等待另一个单元，从而降低整体系统吞吐量。
    • 缓解措施：确保已启用流水线处理（例如，tf_xla_disable_full_embedding_pipelining = false 或其等效项）。
  • 限制导致的瓶颈：
    • 问题：
      • 限制过低：可能会触发过多的微批处理（将输入批次拆分为许多较小的子批次以满足严格的限制）。虽然这样可以保持正确性，但每个小批次都会带来一些处理开销，可能会减慢整体执行速度。如果 allow_id_dropping 为 true，过低的限制也可能导致 ID 丢弃，从而影响模型准确性。
      • 限制过高（但仍适合）：虽然过高的限制可能会阻止小批量处理，但如果实际数据特征很少接近这些峰值，则可能会不必要地增加 SPMEM 压力。它们还可能导致 HBM 堆栈使用量超出实际需要。
      • 编译失败：如果配置的限制需要比可用物理内存更多的 SPMEM 或 HBM 堆栈，则编译将失败。
    • 缓解措施：确保限额设置正确。
  • 数据分布倾斜：
    • 问题：如果某些 SparseCore 分区始终比其他分区接收到不成比例的更多 ID（表明 ID 分布不佳），则这些过载的 SparseCore 将成为性能瓶颈。
    • 缓解措施：在小批量处理过程中进行 ID 随机化有助于缓解堆叠表（尤其是包含“热门”用户表的堆叠表）的这一问题。仔细分析 ID 分布，以设置适当且均衡的单表限制。
  • 表格堆叠问题：
    • 问题：
      • 堆叠的表格过少：可能不足以有效隐藏 ICI 延迟时间或充分减少处理开销。
      • 堆叠的表过多：可能会导致创建非常大的逻辑表，难以管理或可能超出可用资源限制。
    • 缓解措施：
      • 确保堆叠的表格数量达到最佳值。一般准则建议堆叠 5-100 个表，以达到最佳效果。
  • 低效的数值/量化：
    • 问题：如果使用 BF16 或量化整数等较低精度格式就足够了（并且可以提供更快的计算速度），那么使用完整的 FP32 精度可能会成为性能瓶颈。
    • 缓解措施：探索精度较低的选项。不过，请注意，量化本身会产生一些开销，可能需要仔细调整量化参数才能保持模型准确率。
  • HBM 带宽饱和度：
    • 问题：HBM 的数据传入和传出过多，这可能是由极小的特征宽度（导致填充开销过高）、低效的内存访问模式或极多的查找次数造成的，可能会使可用的 HBM 带宽饱和。
    • 缓解措施：增加 TPU 数量有助于缓解 HBM 带宽饱和问题。

常见的内存瓶颈：

  • SPMEM 溢出（编译失败）：
    • 问题：如果 max_ids_per_partition 和 max_unique_ids_per_partition 设置得过高，XLA 编译器可能无法分配足够的 SPMEM，从而导致编译错误，例如："Fixed size allocations (...) do not fit in TileSpmem (...)"。此外，如果 (sample_count * feature_width) / kNumTiles（其中 kNumTiles 是每个 SC 的 tile 数）对于在 tile SPMEM 中暂存收集操作数而言过大，则可能会出现 "Gather operand too large..." 等错误。
    • 缓解措施：减小批次大小或增加用于处理的芯片数量。
  • HBM 堆栈溢出（运行时或编译时）：
    • 问题：如果 feature_width、max_unique_nz_per_row 和 logical_replica_count 的组合导致 HBM 堆栈内存需求超过可用的 HBM，则可能会在运行时或编译期间导致内存不足 (OOM) 错误。
    • 缓解措施：调整 xla_sc_num_serialized_tables_to_optimize_hbm 标志，通过序列化表处理来减少 HBM 堆栈使用量（这通常会牺牲性能）。
  • HBM 堆耗尽：
    • 问题：主要由以下原因导致：非常大的密集层权重、存储在内存中的大量常量，或过于激进的输入预取（高 maximum_parallel_iterations）。
    • 缓解措施：使用 XProf Memory Viewer 等工具监控堆使用情况。
  • 填充开销：
    • 问题：嵌入表在特征维度上填充为 32B 对齐（相当于 8 个浮点数）。因此，较小的特征宽度（例如 1 个浮点数）会产生大量的填充开销（例如，所分配的缓冲区空间中有 7/8 是填充），从而导致 HBM 浪费。表的词汇维度也会填充为系统中 SparseCore 数量的倍数；不过，对于词汇量足够大的表，这种影响通常可以忽略不计。

影响性能和内存的一般因素：

  • 拓扑：可用芯片的数量及其互连架构。
  • 批次大小：直接影响每个 SparseCore 的 sample_count，进而影响内存消耗和计算负载。
  • 数据格式设置：确保高效的设备端数据布局对于实现最佳性能至关重要。

11. 分析 SparseCore 配置文件

分析性能配置文件是识别 SparseCore 工作负载中的瓶颈并发现优化机会的关键步骤。

  1. 获取轨迹：
    • 利用性能分析工具（例如 XProf）在模型训练或运行推理时捕获详细的执行轨迹。此轨迹将提供主机、TensorCore 和 SparseCore 上发生的操作的时间轴。
  2. 检查 Trace Viewer（例如在 XProf 或 TensorBoard 中）：
    • 主机活动：仔细检查主机的活动。TPU 活动是否存在明显缺口？此类缺口可能表明主机是瓶颈，无法快速馈送数据。分析输入流水线的性能。
    • TensorCore (TC) 和 SparseCore (SC) 活动：
      • 查看 TC 和 SC 的执行时间轴。它们是否并行运行，表明流水线是否有效？或者，是否存在一个单元空闲等待另一个单元的长时间段？
      • 确定 SC 和 TC 上耗时最多的操作（运行时间最长的操作）。
      • 直观的跟踪输出（通常显示彩色块，表示不同操作随时间的变化，如 TPU:0 SparseCore 1 (pid 1005)）对于直观地识别主要操作和空闲时段非常有用。
    • 步进时间分析：观察总体步进时间，了解其在主机处理、SC 计算和 TC 计算之间的分布或细分情况。
  3. 内存分析 (XProf Memory Viewer)：
    • 堆使用情况：使用 XProf 的“内存查看器”标签页等工具检查 HBM 堆使用情况。这有助于确定大型模型权重、常量或过于激进的输入预取是否消耗了过多的 HBM。启用 --vmodule=best_fit_allocator=1 等标志可能会提供峰值堆使用情况的日志。
    • 堆栈使用情况（间接）：虽然直接 HBM 堆栈分析可能很复杂，但如果您遇到内存不足错误，并且堆使用情况看起来合理，那么 HBM 堆栈耗尽（通常是由于限制或功能宽度过大）很有可能是原因。提供的 HBM 堆栈使用率公式有助于估算这一点。
  4. 查找特定模式：
    • 小批次处理：如果经常超出限制，您可能会在轨迹中看到小批次处理的证据（例如，与全局批次大小相比，SC 操作的数量更多，但规模更小）。通常可以从日志中或通过观察某些操作的调用次数来推断出这一点。
    • ID 丢弃：如果已启用 ID 丢弃功能，并且正在发生 ID 丢弃，系统日志可能会提供相关指示。这也清楚地表明，配置的限额对于输入数据而言过于严格。
    • 编译时间：重新编译时间过长，尤其是在启用反馈导向优化 (FDO) 并频繁调整限制的情况下，会显著增加总体训练时间开销。
  5. 与标志和配置相关联：
    • 将配置文件中观察到的行为与 SparseCore 配置（限制文件中的设置、XLA 标志）相关联。例如，如果将 xla_sc_num_serialized_tables_to_optimize_hbm 设置为较高的值，您可能会预期 SC 性能会降低，但 HBM 堆栈消耗会降低。
  6. 迭代过程：
    • 分析通常是一个迭代优化过程。进行特定更改（调整限额、启用或停用某项功能），捕获新配置文件，然后将其与之前的配置文件进行比较，以了解修改的影响。

12. 常规调试标志

您可以启用多个标志来帮助调试与 SparseCore 执行相关的问题。请务必注意，启用这些检查通常会降低性能，因此，在生产运行中，这些检查通常应处于停用状态。

  • ID 检查（超出范围）：
    • 标志：xla_sparse_core_enable_id_bound_check = true
    • 用途：在宿主系统上启用检查，以检测输入数据中的任何嵌入 ID 是否超出为给定嵌入表定义的有效词汇范围。这有助于发现与输入数据不正确或损坏相关的问题。
  • NaN 检查器：
    • 标志：xla_sc_detect_nan = true
    • 用途：用于检测在 SparseCore 上处理的浮点数据中的 NaN（非数字）值。如果在各种编译器传递的输入或输出中检测到 NaN，此标志将导致引发错误。此类错误通常会提供有关遇到 NaN 的位置的信息。
  • 边界检查器（内存访问）：
    • 标志：xla_sc_assert_level=bounds
    • 用途：此标志可启用 ASAN (AddressSanitizer) 风格的工具，该工具会重写内存访问指令（例如 VMEM 加载/存储和 DMA 操作），以包含动态检查。这些检查会验证内存访问是否在目标内存区域的分配范围内。
    • 行为：如果检测到越界内存访问，执行将失败。
    • 注意：此检查器可能会产生误报，例如，由于检查器无法完全理解复杂的跨步访问模式。此转换在后端编译过程的后期阶段应用。
  • 缓冲区检查器（内存损坏）：
    • 标志：
      • xla_tpu_buffer_contents_sanitizer_config='cores_to_sanitize: [TC, SC_SCS, SC_TILE], sanitizer_mode: LOCAL_ONLY'
      • xla_tpu_verify_launch_id_across_cores=true
    • 用途：这些标志有助于确保内存缓冲区不会被不相关的操作意外损坏或覆盖。缓冲区清理器会检查缓冲区的内容，以验证它们是否未发生意外变化。

13. 量化支持

SparseCore 的 SparseDenseMatmulOp 旨在支持使用 32 位浮点 (FP32) 和整数数据类型对嵌入表执行操作。虽然模型训练通常使用 FP32 精度进行嵌入表训练，但也可以应用训练后量化 (PTQ)。PTQ 允许在推理中使用较低精度的数据类型（例如 8 位整数），这可能会提高性能并减少内存占用空间。

模拟量化：

SparseDenseMatmulOp 可配置为执行“模拟量化”。在此运行模式下，嵌入向量会先量化为较低精度，然后再反量化为较高精度（例如 FP32），之后才能用于后续计算。此技术可让模型在训练时考虑量化噪声的影响。使用模拟量化进行训练可以在模型完全量化以进行推理时提高最终模型的准确性。

SparseDenseMatmulOp 的配置属性（用于量化）：

  • quantization_config_num_buckets = 256
    • 此属性用于指定将 32 位浮点数量化为的离散分桶或级别数。例如，在量化为 8 位整数时，通常会指定 2^8 =256 个分桶。
  • quantization_config_low = -X.X
    • 此属性用于定义量化范围内的最小浮点值。在量化期间，任何低于此指定最小值的所有输入值都会被剪裁为此最小值。
  • quantization_config_high = Y.Y
    • 此属性定义了量化范围内的最大浮点值。在量化期间，任何高于此指定最大值的输入值都会被剪裁为此最大值。

数值和流水线交互：

模型的数值行为可能会因 TensorCore 和 SparseCore 之间的流水线是否已启用而异。如果流水线处于活跃状态，SparseCore 处理的梯度可能“过时”（来自上一次迭代）。这可能会与量化过程互动，并可能影响模型训练动态或最终准确率。

14. 即将推出的功能和近期改进

SparseCore 生态系统会不断开发和增强。

路线图：

  • 样本维度的小批次处理：
    • 此功能计划作为对现有词汇维度小批处理功能的补充。
    • 这样就可以沿样本维度进一步划分嵌入输入。这可以通过引入设备端循环来实现，该循环可以一次过滤和处理来自一部分样本的查找。此功能有助于管理非常大的每个样本 ID 数量，或改进处理单元之间的负载平衡。
  • 改进了对每行少于 8 个整数的嵌入（特征宽度较小）的支持：
    • 当前设计通常会为宽度小于 8 个浮点数（相当于 32 字节）的嵌入特征使用大量填充。这种填充可能会导致 HBM 浪费，并可能导致计算资源利用不足。未来的改进旨在缓解小特征维度表的这种低效问题。

近期改进：

  • 在 HBM 中暂存收集操作数：
    • 此优化有助于通过允许在较大的 HBM 中暂存一些收集操作输入或输出来减轻共享暂存存储器 (SPMEM) 的压力。
  • 堆栈内存用量减少：
    • 我们已实施多项增强功能，以减少 SparseCore 运行期间的 HBM 堆栈内存消耗，理想情况下不会对整体性能或吞吐量产生负面影响。

这些增强功能旨在提高 SparseCore 的性能、内存效率和操作灵活性，以适应更广泛的稀疏工作负载。