GCP使用Tensorflow-cloud训练TPU模型chief_config报错排查

问题背景 #

在Google Cloud Platform上训练搭载CNN分类头的Longformer-large大模型时，采用Tensorflow-cloud搭配Colab提交训练任务。此前使用batch size=4、4张P100 GPU的配置仍触发OOM错误，因此计划改用TPU资源，将batch size提升至8后运行报错，提示TPU配置不能作为chief_worker_config。

错误复现 #

错误调用代码 #

tfc.run(
distribution_strategy="auto",
requirements_txt="requirements.txt",
docker_config=tfc.DockerConfig(
   image_build_bucket=GCS_BUCKET
),
worker_count=1,
worker_config= tfc.COMMON_MACHINE_CONFIGS["TPU"],
chief_config=tfc.COMMON_MACHINE_CONFIGS["TPU"],
job_labels={"job": JOB_NAME})

核心报错信息 #

ValueError: Invalid `chief_config` input. `chief_config` cannot be a TPU config. Received <tensorflow_cloud.core.machine_config.MachineConfig object at 0x7f5860afe210>.

正确配置方案 #

方案1：TPU资源正确配置 #

报错的核心原因是Tensorflow-cloud的集群规则中，chief节点负责调度任务、汇总日志、存储checkpoint，必须使用CPU机型，不能配置为TPU资源，TPU仅能作为计算worker节点挂载。
修正后的调用代码如下：

tfc.run(
    distribution_strategy="auto",
    requirements_txt="requirements.txt",
    docker_config=tfc.DockerConfig(
        image_build_bucket=GCS_BUCKET
    ),
    chief_config=tfc.COMMON_MACHINE_CONFIGS["CPU"], # chief节点固定用CPU配置
    worker_count=1,
    worker_config=tfc.COMMON_MACHINE_CONFIGS["TPU"],
    job_labels={"job": JOB_NAME}
)

使用TPU时的注意事项：

• 训练数据集、模型保存路径必须全部挂载到GCS存储桶，TPU节点无法读取Colab本地文件
• 确认所有模型算子支持TPU编译：Longformer主干、标准Keras实现的CNN分类头均无兼容问题
• 单TPU v3-8节点的HBM容量足够支撑batch size=8的Longformer-large训练，不会触发OOM

方案2：GPU资源稳定配置（无TPU兼容成本） #

如果不想处理TPU的文件访问、算子兼容问题，仅需在原有GPU训练逻辑基础上做两处显存优化，不需要高配4卡P100即可稳定运行：

• 开启混合精度训练，在训练代码入口添加tf.keras.mixed_precision.set_global_policy('mixed_float16')，可直接降低50%左右显存占用
• 给Longformer主干开启梯度检查点，调用longformer_backbone.gradient_checkpointing_enable()，可再降低30%-40%显存占用

优化后单张A100 40G即可稳定跑batch size=8的训练任务，对应tfc配置代码如下：

tfc.run(
    distribution_strategy="tf.distribute.MirroredStrategy",
    requirements_txt="requirements.txt",
    docker_config=tfc.DockerConfig(
        image_build_bucket=GCS_BUCKET
    ),
    chief_config=tfc.COMMON_MACHINE_CONFIGS["CPU"],
    worker_count=1,
    worker_config=tfc.COMMON_MACHINE_CONFIGS["A100_1X"],
    job_labels={"job": JOB_NAME}
)

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