在TPU VM上训练T5X模型触发核心转储终止

我在针对Winograd Schema Challenge训练T5X模型时，运行训练脚本出现以下错误：

2022-08-21 17:27:01.141608: F ./tensorflow/core/tpu/tpu_library_init_fns.inc:101] TpuEmbeddingEngine_ConfigureCommunication not available in this library.
Aborted (core dumped) 

错误原因分析 #

• TPU库与TensorFlow版本不兼容：这个错误核心是当前使用的TensorFlow版本缺少TpuEmbeddingEngine_ConfigureCommunication函数，T5X对TF和TPU runtime的版本匹配要求极高，比如早期T5X版本仅支持特定分支的TensorFlow 2.x或预编译的TPU专属TF包，版本不匹配就会出现这类函数缺失问题。
• TPU环境初始化不完整：如果使用云TPU，可能未正确设置TPU节点的环境变量，或是使用了不匹配的TPU驱动镜像；如果是本地模拟TPU，可能未安装对应版本的TPU模拟器依赖。
• T5X训练配置有误：训练配置文件中可能启用了TPU嵌入功能（比如use_tpu_embedding设为True），但当前环境的硬件或TF版本并不支持该特性。

内容的提问来源于stack exchange，提问作者Akash Kumar