使用TPU训练RoBERTa模型时预测耗时远超训练的问题排查

我在Google Colab中基于自定义数据集从零训练RoBERTa模型，用TPU后训练耗时大幅缩短（CPU需数小时，TPU仅需数分钟），但模型预测耗时却远超训练：预测数据量仅为训练的1/4，耗时却超过训练的两倍，相当于训练速度比预测快10倍。我本来预期预测更快，已经试过修改输入格式（numpy、张量）、调整batch size，但都没改善。训练速度正常说明TPU配置没问题，预测用的也是同一份数据集，数据本身没毛病，想知道为啥预测这么慢。

代码细节 #

导入依赖 #

from transformers import RobertaConfig
from transformers import RobertaForMaskedLM
from transformers import TFRobertaForMaskedLM
import tensorflow as tf

环境警告 #

/usr/local/lib/python3.11/dist-packages/torch_xla/__init__.py:253: UserWarning: `tensorflow` can conflict with `torch-xla`. Prefer `tensorflow-cpu` when using PyTorch/XLA. To silence this warning, `pip uninstall -y tensorflow && pip install tensorflow-cpu`. If you are in a notebook environment such as Colab or Kaggle, restart your notebook runtime afterwards.
warnings.warn(

TPU配置 #

resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)
print("replicas",strategy.num_replicas_in_sync)
batch_size = 8 * strategy.num_replicas_in_sync
print(strategy)

配置输出

replicas 8
<tensorflow.python.distribute.tpu_strategy.TPUStrategyV2 object at 0x7bdcd1c6d6d0>

模型初始化 #

config = RobertaConfig(
    vocab_size=10_000,
    max_position_embeddings=64,
    num_attention_heads=6,
    num_hidden_layers=3,
    type_vocab_size=1,
    hidden_size=300,
    intermediate_size=600
)
with strategy.scope():
  model = TFRobertaForMaskedLM(config=config) #Tensorflow Roberta Model
  model.compile(
      optimizer=tf.keras.optimizers.AdamW(learning_rate=5e-05)
  )

数据加载 #

import tensorflow as tf
def decode_fn(sample):
    features = {
        "input_ids": tf.io.FixedLenFeature((64,), dtype=tf.int64),
        "attention_mask": tf.io.FixedLenFeature((64,), dtype=tf.int64),
        "labels": tf.io.FixedLenFeature((64,), dtype=tf.int64)
    }
    return tf.io.parse_example(sample, features)
tf_dataset = tf.data.TFRecordDataset(["dataset.tfrecords"])
tf_dataset = tf_dataset.map(decode_fn)
tf_dataset = tf_dataset.batch(batch_size, drop_remainder=True)
tf_dataset = tf_dataset.apply(
    tf.data.experimental.assert_cardinality(263317 // 64))
predict_check_dataset = tf_dataset.take(1000)

数据批次统计

print("Number of training Batch:",len(list(tf_dataset)))
print("Number of prediction Batch:",len(list(predict_check_dataset)))

统计输出

Number of training Batch: 4114
Number of prediction Batch: 1000

训练过程 #

train_log = model.fit(tf_dataset)

训练日志

4114/4114 [==============================] - 94s 23ms/step - loss: 1.5873

预测过程 #

for batch in predict_check_dataset:
        prediction = model(batch['input_ids'], attention_mask=batch['attention_mask'],training=False)

预测耗时约240秒，而训练仅耗时94秒，尽管预测数据量仅为训练的1/4。

问题原因及解决方案 #

核心原因 #

1. 预测未利用TPU分布式并行：model.fit()会自动把批次拆分到8个TPU核心并行处理，但手动遍历数据集调用model()的方式，完全没用到TPU的分布式能力，每个批次都只在单核心运行。
2. 数据流水线未优化：训练时TF Data会自动做预取、缓存等优化，但预测数据集没复用这些逻辑，手动循环还会引入Python层面的调度开销。
3. Eager模式效率低：直接调用model()属于eager执行模式，在TPU上的调度开销远大于训练时的graph执行模式。

解决步骤 #

• 改用分布式预测接口：用model.predict()替代手动循环，它会自动适配TPU策略，利用多核心并行：

  # 自动利用TPU分布式的预测方式
  predictions = model.predict(predict_check_dataset, batch_size=batch_size)
  

• 优化预测数据流水线：给预测数据集加上缓存和预取，减少IO等待：

  predict_check_dataset = predict_check_dataset.cache().prefetch(tf.data.AUTOTUNE)
  

• 处理环境冲突：根据警告提示卸载tensorflow，安装tensorflow-cpu，避免torch-xla和tensorflow抢占资源干扰TPU调度。

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