在TPUv4-64上保存Orbax checkpoint时遭遇打开文件过多错误

原因分析 #

这个错误并非单纯的单节点文件句柄不足——TPU Pod是多节点分布式集群，仅调大本地节点的ulimit无法覆盖所有节点；同时Orbax默认的分布式 checkpoint 策略在GCS存储下可能会打开过多临时对象句柄，导致全局句柄耗尽。

具体解决方案 #

1. 全局配置所有TPU节点的文件句柄限制 #

仅在本地执行ulimit -n 65536无法作用于整个TPU Pod，需要在集群启动阶段为所有节点设置上限：

# 在TPU Pod的全局启动脚本中执行
sudo sysctl -w fs.file-max=1000000
echo "* soft nofile 65536" >> /etc/security/limits.conf
echo "* hard nofile 65536" >> /etc/security/limits.conf

如果使用Vertex AI TPU训练，需将上述配置加入自定义容器镜像或启动命令中，确保所有TPU节点生效。

2. 调整Orbax的Checkpoint参数 #

修改Orbax的并行IO策略，减少文件句柄占用：

from orbax.checkpoint import CheckpointManager, CheckpointManagerOptions
from orbax.checkpoint import PyTreeCheckpointer

# 降低并行IO数量，增大缓冲区减少IO频次
options = CheckpointManagerOptions(
    max_to_keep=5,
    file_io_options={
        'max_parallelism': 16,
        'buffer_size': 1024 * 1024 * 64,
    }
)
checkpointer = PyTreeCheckpointer()
manager = CheckpointManager(
    '/path/to/checkpoint-test',
    checkpointer,
    options=options,
)

同时可以设置max_shard_size参数，让每个checkpoint文件更大，减少总文件数量，进一步降低句柄占用。

3. 清理GCS临时文件并优化临时目录策略 #

Orbax checkpoint过程中会生成大量临时文件，分布式场景下可能出现句柄泄漏：

• 训练前手动清理旧的临时文件：

gsutil rm -r /path/to/checkpoint-test/*orbax-checkpoint-tmp*

• 为每个TPU节点配置独立的临时目录前缀，避免跨节点的临时文件冲突。

4. 验证TPU节点的实际句柄限制 #

登录每个TPU节点，确认文件句柄限制是否生效：

cat /proc/sys/fs/file-max
ulimit -n

如果file-max仍然偏低，需通过sysctl永久修改，或联系GCP支持调整TPU VM的默认配置。

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