Apple M1设备上TensorFlow+BERT模型GPU运行报错求助

问题背景 #

在Apple M1 Mac上搭建基于BERT的文本分类器，采用TensorFlow 2.9.0、tf-text 2.9.0、tf-metal 0.5.0版本，模型代码如下：

text_input = tf.keras.layers.Input(shape=(), dtype=tf.string, name = "text") 
encoder_inputs = bert_preprocess(text_input) 
encoder_outputs = bert_encoder(encoder_inputs) 
bert_embeds = encoder_outputs["pooled_output"] 
intermediate_layer = tf.keras.layers.Dense(512, activation = "relu", name = "intermediate_layer")(bert_embeds) 
dropout_layer = tf.keras.layers.Dropout(0.1, name = "dropout_layer")(intermediate_layer) 
output_layer = tf.keras.layers.Dense(1, activation = "sigmoid", name = "output_layer")(dropout_layer) 
model = tf.keras.Model(text_input, output_layer)

训练时出现以下GPU兼容性错误：

2 root error(s) found. (0) NOT_FOUND: No registered 'AddN' OpKernel for 'GPU' devices compatible with node {{node model_13/keras_layer_1/StatefulPartitionedCall/StatefulPartitionedCall/StatefulPartitionedCall/bert_pack_inputs/PartitionedCall/RaggedConcat/ArithmeticOptimizer/AddOpsRewrite_Leaf_0_add_2}} (OpKernel was found, but attributes didn't match) Requested Attributes: N=2, T=DT_INT64, _XlaHasReferenceVars=false, _grappler_ArithmeticOptimizer_AddOpsRewriteStage=true, _device="/job:localhost/replica:0/task:0/device:GPU:0" . Registered: device='XLA_CPU_JIT'; T in [DT_FLOAT, DT_DOUBLE, DT_INT32, DT_UINT8, DT_INT16, 16534343205130372495, DT_COMPLEX128, DT_HALF, DT_UINT32, DT_UINT64, DT_VARIANT] device='GPU'; T in [DT_FLOAT] device='DEFAULT'; T in [DT_INT32] device='CPU'; T in [DT_UINT64] device='CPU'; T in [DT_INT64] device='CPU'; T in [DT_UINT32] device='CPU'; T in [DT_UINT16] device='CPU'; T in [DT_INT16] device='CPU'; T in [DT_UINT8] device='CPU'; T in [DT_INT8] device='CPU'; T in [DT_INT32] device='CPU'; T in [DT_HALF] device='CPU'; T in [DT_BFLOAT16] device='CPU'; T in [DT_FLOAT] device='CPU'; T in [DT_DOUBLE] device='CPU'; T in [DT_COMPLEX64] device='CPU'; T in [DT_COMPLEX128] device='CPU'; T in [DT_VARIANT]

模型在CPU上可正常运行但速度极慢，需实现GPU加速运行。

错误分析 #

错误核心是：tf-metal的GPU版本AddN OpKernel仅支持DT_FLOAT类型，但当前BERT预处理阶段的某个节点请求了DT_INT64类型的AddN操作。这是因为BERT预处理包含大量整数型序列操作（如token ID拼接、序列长度计算），而旧版本tf-metal对整数类型的GPU运算支持不完善，无法处理int64的GPU加法操作。

解决方案 #

1. 强制预处理操作在CPU执行 #

将BERT预处理层限定在CPU运行，模型主体（BERT编码器、全连接层等浮点运算部分）仍在GPU执行，避开GPU对整数操作的支持短板：

# 预处理层绑定CPU设备
with tf.device('/CPU:0'):
    text_input = tf.keras.layers.Input(shape=(), dtype=tf.string, name="text")
    encoder_inputs = bert_preprocess(text_input)

# 模型主体自动调度到GPU
encoder_outputs = bert_encoder(encoder_inputs)
bert_embeds = encoder_outputs["pooled_output"]
intermediate_layer = tf.keras.layers.Dense(512, activation="relu", name="intermediate_layer")(bert_embeds)
dropout_layer = tf.keras.layers.Dropout(0.1, name="dropout_layer")(intermediate_layer)
output_layer = tf.keras.layers.Dense(1, activation="sigmoid", name="output_layer")(dropout_layer)
model = tf.keras.Model(text_input, output_layer)

2. 升级tf-metal版本 #

tf-metal 0.5.0对整数类型的GPU支持有限，升级到兼容TensorFlow 2.9的更高版本（如0.6.0）可改善该问题：

pip install --upgrade tensorflow-metal==0.6.0

注：需确认tf-metal版本与TensorFlow版本的兼容性，可参考官方兼容矩阵。

3. 禁用XLA优化 #

XLA优化可能会强制将整数操作调度到GPU，尝试关闭XLA编译：

# 全局禁用XLA
tf.config.optimizer.set_jit(False)

# 或在模型编译时关闭
model.compile(
    optimizer='adam',
    loss='binary_crossentropy',
    metrics=['accuracy'],
    jit_compile=False
)

4. 调整数据类型为int32 #

tf-metal的GPU对int32的支持优于int64，将预处理阶段的整数操作类型改为int32：

# 数据加载时转换标签类型为int32
train_dataset = train_dataset.map(lambda x, y: (x, tf.cast(y, tf.int32)))

# 若使用自定义预处理函数，确保序列长度等参数为int32
def preprocess_fn(text, label):
    processed = bert_preprocess(text)
    processed['input_ids'] = tf.cast(processed['input_ids'], tf.int32)
    processed['attention_mask'] = tf.cast(processed['attention_mask'], tf.int32)
    return processed, tf.cast(label, tf.int32)

验证方法 #

应用上述方案后，重新编译模型并启动训练，通过以下方式确认GPU是否正常运行：

• 打开Activity Monitor，切换到GPU标签，查看TensorFlow进程的GPU使用率
• 在代码中添加设备检测：

print("GPU可用:", tf.config.list_physical_devices('GPU'))

内容的提问来源于stack exchange，提问作者amin-oueslati