TensorFlow 中 TPU 训练实战：从 GPU 迁移到 TPU 的完整指南

1. 为什么要在 TensorFlow 里认真考虑 TPU 这条路线很多人第一次接触 TPU脑子里冒出来的第一个念头是这玩意儿是不是只有大厂才用得起。我一开始也这么想直到自己真正把一个中等规模的图像分类模型从单卡 GPU 迁到 TPU 上跑了一遍才发现事情没那么玄乎。TPU 全称 Tensor Processing Unit是专门为张量运算设计的加速芯片它和 GPU 最大的区别在于GPU 是通用并行计算器什么活都能干一点TPU 更像是一台矩阵乘法专用机床只干一件事但干得极快、极省电。TensorFlow 和 TPU 的关系说白了就是亲儿子级别的适配。TPU 从设计之初就是围绕 TensorFlow 的计算图来做的所以你在 TensorFlow 里调用 TPU比在其他框架里调用要顺滑得多。这也是为什么很多做大规模训练的团队即便平时用别的框架做实验到了要压榨算力的时候还是会回到 TensorFlow TPU 这套组合上。那到底什么场景适合上 TPU我总结下来是三类第一类是模型参数量大、矩阵运算密集的比如 Transformer 系列、大型卷积网络第二类是 batch size 能开得很大的因为 TPU 的算力要靠大 batch 才能喂饱第三类是你需要长时间、多轮次训练对单位算力成本敏感的。反过来如果你的模型很小、batch size 只有几十、还频繁做动态控制流那 TPU 反而可能不如 GPU 来得舒服。这篇文章我打算按真实迁移的顺序来讲先搞清楚 TPU 的硬件形态和访问方式再讲清楚 TensorFlow 里那套TPUStrategy到底怎么用然后是数据管道怎么改、模型代码要注意什么最后聊聊我踩过的那些坑。全程围绕 TensorFlow 生态不跑偏。提示TPU 不是插上就能用的显卡它是一台独立的计算资源你需要通过网络去访问它。这个心智模型如果不建立起来后面很多设计你会想不通。2. TPU 的几种形态与访问方式先弄明白你在跟谁打交道2.1 从本地插卡到远程计算节点的认知转变GPU 的用法大家都熟插在机器上nvidia-smi一看/gpu:0直接就能用。TPU 完全不是这个逻辑。你本地那台笔记本上基本不可能有 TPU它是部署在远程计算集群里的独立加速器。你的训练脚本跑在一台宿主虚拟机上通过特定的网络地址去连接 TPU 工作节点计算图被编译后下发到 TPU 上执行。这个架构带来的第一个直接影响就是你的代码必须显式地声明我要连哪个 TPU。在 TensorFlow 里这个声明通常通过环境变量或者初始化函数来完成。如果你用的是托管式的 TPU 资源平台会给你一个形如grpc://10.x.x.x:8470的地址你的脚本要拿着这个地址去建立连接。2.2 单核、多核与 Pod算力规模的三个档位TPU 的算力不是一块卡的概念而是按核来算的。一个 TPU 芯片里通常有多个核心每个核心有独立的矩阵运算单元和内存。常见的形态有这么几种形态核心数量级典型用途通信方式单核 TPU1 核小模型调试、跑通流程无跨核通信多核 TPU4~8 核中等规模训练芯片内高速互联TPU Pod数十到上千核超大模型、大规模分布式专用高速网络对绝大多数个人开发者和中小团队来说能接触到的基本是前两种。Pod 级别的资源通常需要排队申请而且对代码的分布式写法要求更高。我建议你先把单核和多核跑顺别一上来就盯着 Pod否则调试成本会让你怀疑人生。2.3 连接 TPU 的两种典型姿势在 TensorFlow 里连接 TPU常见的有两种方式。一种是手动指定地址用tf.distribute.cluster_resolver.TPUClusterResolver去解析另一种是在托管环境里直接调用平台提供的初始化函数让它自动发现资源。前者更通用后者更省事。import tensorflow as tf # 方式一手动指定 TPU 地址 resolver tf.distribute.cluster_resolver.TPUClusterResolver(tpugrpc://10.0.0.1:8470) tf.config.experimental_connect_to_cluster(resolver) tf.tpu.experimental.initialize_tpu_system(resolver) # 方式二托管环境自动发现不同平台函数名可能不同 # resolver tf.distribute.cluster_resolver.TPUClusterResolver() # tf.config.experimental_connect_to_cluster(resolver) # tf.tpu.experimental.initialize_tpu_system(resolver)这里有个细节很多人会忽略initialize_tpu_system这一步是必须的而且要在创建分布策略之前调用。它的作用是初始化 TPU 设备、加载必要的运行时。如果你跳过这一步直接建策略会报设备找不到的错。我见过不止一个同学卡在这里以为是权限问题其实是漏了初始化。注意initialize_tpu_system在同一个进程里重复调用可能出问题建议用一次性的初始化逻辑包起来别在循环里反复调。3. TPUStrategy 的核心用法把训练逻辑翻译给 TPU 听3.1 分布策略是 TensorFlow 处理多设备的统一入口TensorFlow 从 2.x 开始把在多个设备上跑同一份计算这件事抽象成了tf.distribute.Strategy。GPU 多卡用MirroredStrategyTPU 就用TPUStrategy。这个抽象的好处是你的模型定义、损失函数、训练循环写法基本不用大改只要把创建变量和训练步骤包在策略的scope里就行。strategy tf.distribute.TPUStrategy(resolver) with strategy.scope(): model build_model() model.compile( optimizertf.keras.optimizers.Adam(learning_rate1e-3), losstf.keras.losses.SparseCategoricalCrossentropy(from_logitsTrue), metrics[accuracy] )关键点在于with strategy.scope():这个上下文。所有在这个 scope 里创建的变量都会被策略接管分配到各个 TPU 核心上并且自动处理跨核的同步。如果你在 scope 外面建了模型那模型变量就落在宿主机的 CPU 上TPU 根本用不上训练速度会慢到让你以为没生效。3.2 全局 batch size 与每核 batch size 的关系这是 TPU 训练里最容易算错的一个数。假设你有 8 个 TPU 核心你想让每个核心每次处理 128 个样本那么你传给model.fit的 batch size 应该是 8 × 128 1024而不是 128。TPUStrategy 会自动把这个全局 batch 切分到各个核心上。我踩过的坑是这样的一开始我按 GPU 的习惯写了 batch size 64结果 8 个核心每个只分到 8 个样本算力严重浪费训练慢得离谱。后来把全局 batch 提到 1024速度直接上了一个台阶。所以你在调 batch size 的时候心里要有一本账全局 batch 每核 batch × 核心数。当然batch 也不是越大越好。太大的 batch 会影响收敛可能需要配合学习率 warmup 或者调整优化器。我的经验是先把每核 batch 定在 128 左右再根据显存TPU 叫 HBM和收敛情况微调。3.3 在策略下评估和预测也要保持一致很多人训练部分包了 scope评估和预测却忘了。虽然评估不一定报错但为了行为一致建议评估也走同样的策略。model.evaluate和model.predict在 TPUStrategy 下会自动做分布式处理你只要保证模型是在 scope 里建的就行。# 评估同样受益于分布式 loss, acc model.evaluate(eval_dataset, stepseval_steps) print(feval loss: {loss:.4f}, eval acc: {acc:.4f})这里steps参数在 TPU 上要特别注意。因为 TPU 是静态图编译执行它需要知道总共要跑多少步。如果你的数据集是tf.data的无限流必须显式给steps否则会一直跑下去或者报错。4. 数据管道改造TPU 最怕喂不饱4.1 为什么 TPU 对输入管道比 GPU 更挑剔GPU 训练时如果数据加载慢GPU 利用率掉一点你还能忍。TPU 不行。TPU 的算力太集中了一旦输入管道跟不上核心就会空转而且这种空转在分布式环境下会被放大。更麻烦的是TPU 的编译执行模式对动态性容忍度低如果你的tf.data管道里有大量 Python 侧的逻辑性能会断崖式下跌。所以给 TPU 准备数据核心原则就一条把能放到图里执行的都放进去把能预取的都预取。tf.data的map、batch、prefetch这些操作尽量用 TensorFlow 原生算子少用tf.py_function。4.2 一个能跑顺的输入管道长什么样我常用的模板大概是这样def build_dataset(filenames, is_training, batch_size): dataset tf.data.TFRecordDataset(filenames, num_parallel_readstf.data.AUTOTUNE) def parse_fn(record): features tf.io.parse_single_example(record, feature_description) image tf.io.decode_jpeg(features[image], channels3) image tf.image.resize(image, [224, 224]) image tf.cast(image, tf.float32) / 255.0 label features[label] return image, label dataset dataset.map(parse_fn, num_parallel_callstf.data.AUTOTUNE) if is_training: dataset dataset.shuffle(10000) dataset dataset.repeat() dataset dataset.batch(batch_size, drop_remainderTrue) dataset dataset.prefetch(tf.data.AUTOTUNE) return dataset几个关键点值得展开说。num_parallel_reads和num_parallel_calls都设成AUTOTUNE让 TensorFlow 自己根据 CPU 核数决定并行度比手动写死数字靠谱。drop_remainderTrue在 TPU 上几乎是必须的因为 TPU 的静态形状要求每个 batch 大小一致最后那个不满的 batch 会引发重新编译拖慢整体速度。4.3 预取深度和宿主机内存的平衡prefetch的缓冲区大小也是门学问。设太小预取不够TPU 还是会等设太大宿主机的内存会被吃光尤其是大图像数据集。我一般先用AUTOTUNE如果发现宿主内存吃紧再手动调成一个固定值比如 4 到 8 个 batch。还有一个容易被忽视的点宿主机的 CPU 核数。TPU 的宿主虚拟机通常 CPU 核数不少但如果你在map里做了很重的图像增强CPU 可能成为瓶颈。这时候可以考虑把增强逻辑用 TensorFlow 的tf.image算子重写或者提前把增强后的数据落盘。我做过一次对比把 Python 版的随机裁剪换成tf.image版本后输入吞吐提升了将近一倍。提示判断输入管道是否是瓶颈可以单独跑一个只读数据、不做训练的循环看每秒能产出多少 batch再和训练时的每秒步数对比。如果两者接近说明瓶颈在数据侧。5. 模型代码里那些GPU 能跑、TPU 会炸的写法5.1 动态形状和控制流是 TPU 的天敌TPU 的计算图是提前编译好的它喜欢形状固定、控制流静态的模型。如果你的模型里有依赖运行时张量值的if分支或者形状随输入变化的操作TPU 编译时就会报错或者退化成极慢的执行模式。常见的雷区包括用tf.cond但条件依赖张量值、在模型里做动态 reshape、用 Python 的if判断张量内容。这些在 GPU 上可能只是慢一点在 TPU 上直接就是编译失败。解决办法是尽量把逻辑改成静态的比如用tf.where代替条件分支用固定的 padding 代替动态形状。5.2 自定义训练循环在 TPU 上的注意事项如果你不用model.fit而是自己写tf.GradientTape的训练循环那在 TPU 上要额外小心。TPUStrategy 提供了strategy.run来执行每一步你需要把单步训练逻辑包进去。tf.function def train_step(iterator): def step_fn(inputs): images, labels inputs with tf.GradientTape() as tape: logits model(images, trainingTrue) loss loss_fn(labels, logits) grads tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) return loss per_replica_losses strategy.run(step_fn, args(next(iterator),)) return strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses, axisNone)这里strategy.run会把step_fn分发到各个核心strategy.reduce把各核的损失汇总。注意tf.function装饰器不能少TPU 需要它来做图编译。我第一次写自定义循环时忘了加结果每步都在做 eager 执行速度惨不忍睹。5.3 变量初始化和检查点的坑在 TPU 上保存和加载检查点路径要选在宿主机能访问、且所有核心都能访问的位置。如果你用的是托管环境通常会有指定的存储桶路径。我遇到过检查点保存成功但加载失败的情况排查半天发现是路径权限问题不同核心对路径的可见性不一致。另外tf.train.Checkpoint在 TPUStrategy 下会自动处理跨核的变量同步你不需要手动做 all-reduce。但如果你自己管理变量就要留意每个核心上的变量副本是否一致。6. 实测中那些让我熬夜的报错与排查思路6.1 Cannot find TPU device先查连接再查初始化这个报错出现频率极高。排查顺序我建议是这样第一确认 TPU 地址是否正确能不能 ping 通第二确认initialize_tpu_system是否在创建策略前调用了第三确认当前进程有没有权限访问该资源。我遇到过一次是地址写错了一位数字折腾了半小时。6.2 编译时间过长多半是图太复杂或形状不稳定TPU 第一次执行某个计算图时要编译编译时间从几十秒到几分钟不等。如果你发现每次训练启动都要等很久可能是图里有动态形状导致反复编译。检查方法是看日志里有没有 Compiling computation 反复出现。解决办法是固定输入形状、减少动态操作。6.3 训练 loss 不下降检查学习率和 batch 的匹配从 GPU 迁到 TPU 后全局 batch 变大了如果学习率还按原来的设收敛会变慢甚至不收敛。经验法则是 batch 翻倍学习率也相应调大但要注意配合 warmup。我一般会先用一个较小的学习率跑几百步确认 loss 在降再逐步调大。6.4 内存溢出TPU 的 HBM 也不是无限的TPU 每个核心的内存是有限的模型太大或者 batch 太大都会 OOM。和 GPU 不同的是TPU 的 OOM 报错信息有时候不那么直观。我的做法是先把 batch 调到很小确认模型能跑起来再逐步往上加找到上限。报错现象可能原因排查动作Cannot find TPU device地址错/未初始化/权限检查地址、初始化顺序、访问权限编译时间异常长动态形状/图复杂固定形状、简化控制流loss 不下降学习率与 batch 不匹配调整学习率、加 warmupOOMbatch 或模型过大减小 batch、精简模型输入吞吐低管道瓶颈优化 map、增加并行、调 prefetch7. 迁移完成后我总结的几条实用经验把模型从 GPU 迁到 TPU本质上不是换个设备那么简单而是换了一套执行哲学。GPU 宽容、灵活TPU 严格、高效。你得顺着它的脾气来它才给你好脸色。第一条经验先在单核上把流程跑通再上多核。单核调试成本低报错信息也更清晰。等单核跑顺了改成多核基本就是改个 batch size 的事。第二条经验把 batch size 当成一等公民来对待。它直接决定了 TPU 的利用率和收敛行为值得你花时间反复调。第三条经验数据管道要当成模型的一部分来优化。TPU 算力越强数据侧的重要性越突出。我见过太多人模型调得很精细数据管道却一塌糊涂最后怪 TPU 慢。第四条经验善用日志和性能分析工具。TensorFlow 提供的 profiler 能看到每个操作的耗时帮你定位到底是计算慢还是数据慢。别靠猜靠数据说话。最后分享一个小技巧如果你不确定某段代码在 TPU 上能不能跑可以先在一个极小的输入上试一下看编译能不能过。编译过了再放大规模。这样能省下大量等待时间。TPU 的调试节奏和 GPU 不一样耐心一点把每一步都验证清楚后面会越来越顺。