TensorFlow MirroredStrategy多GPU运行3D U-Net触发GpuLaunchKernel报错

报错原始信息：

GpuLaunchKernel( SwapDimension1And2InTensor3UsingTiles<T, kNumThreads, 
kTileSize, kTileSize, conjugate>, total_tiles_count, kNumThreads, 0, 
d.stream(), input, input_dims, output) status: Internal: 
invalid configuration argument
Aborted (core dumped)

报错翻译：CUDA内核SwapDimension1And2InTensor3UsingTiles启动失败，内部错误：无效的配置参数，进程终止（核心转储）

根因说明 #

该错误属于TensorFlow底层CUDA算子的适配问题，和模型架构、数据生成逻辑无关：

• 单GPU运行时，3D张量维度交换算子的tile尺寸、线程数配置默认适配单卡参数
• 切换到MirroredStrategy多GPU训练后，张量会自动切分到多卡，切分后的子张量尺寸不符合该内核的参数校验规则，导致启动失败

可行修复方案 #

• 调整3D输入的三个维度尺寸，确保所有维度都是16的整数倍，匹配CUDA显存对齐要求，避免维度交换时出现非对齐的张量切片
• 降低全局batch size，保证每张卡分配到的子batch对应的3D数据块尺寸不超过单卡内核运算上限，避免总tile数超出内核支持的最大值
• 升级TensorFlow到2.12及以上版本，该版本修复了多个3D算子在分布式训练下的内核参数校验问题
• 若无法升级TensorFlow版本，替换模型中所有隐式的3D张量维度转置操作，改为tf.transpose显式调用，禁用自动融合的维度交换内核
• 暂时更换分布式策略为MultiWorkerMirroredStrategy验证，排除多卡通信时的张量切片逻辑错误

内容的提问来源于stack exchange，提问作者Yassine Barhoumi