多GPU训练中EagerKernelExecutes后大量DeserializeSparse阶段的原因排查

在AWS g4dn.12xlarge实例的4块GPU上训练一个接收稠密和稀疏张量输入的小型TF2.x模型：仅使用稠密特征时，分布式训练代码运行正常，无性能损耗；但加入稀疏特征后，TensorBoard Profiler的trace_viewer中出现大量异常块。

核心问题 #

• 所有GPU能正常计算分配的批次，但主机侧计算块之间存在较长时间间隔。出现17×4个EagerExecute:DeserializeSparse操作，其终端操作为_Send input 0 from /job:localhost/replica:0/task:0/device:GPU:{gpu_number} to /job:localhost/replica:0/task:0/device:CPU:0（17为模型接收的稀疏特征数量，4为所用GPU数量）。
• 大量MemcpyD2H操作（截图中的粉色小块）占用每个GPU，且未并行执行，该耗时约为实际前向传播的6倍。

模型稀疏张量处理代码 #

def call(self, inputs: tf.sparse.SparseTensor):
  with tf.device("\cpu:0"):
    x = self.hash_inputs_from_static_hash_table(inputs)
    x = self.embedding_lookup_sparse(x)
  return self.prediction_head(x)

数据量并不大（每个副本批次大小为128，稀疏特征嵌入维度<10），已尝试将所有稀疏相关操作移至CPU以减轻GPU负担，但问题仍未解决。

需要明确这些GPU计算后出现异常块的原因，并消除瓶颈以充分利用多GPU分布式训练的优势。

内容的提问来源于stack exchange，提问作者Wonryong Ryou