在训练资源比较紧张的环境里单张 GPU 往往要同时扛好几个实验。我常遇到的情况是白天要跑一个基线模型晚上想再插一个对比实验或者某个大模型把显存占得七七八八剩下的一小块显存既放不下完整 batch又不想让 GPU 完全闲着。TensorFlow 2 在单卡上同时训练多个模型听起来有点“挤牙膏”的意思但真做起来有讲究不是简单把两个 model.fit() 叠在一起就完事。这篇文章就聊我实际试下来比较稳的做法、踩过的坑以及背后的资源分配逻辑给准备在一张卡上同时跑多个模型的同学做个参考。1. 为什么非要在单卡上同时训练资源空转才是最贵的成本先说个很现实的问题一张训练卡如果只跑一个小模型显存利用率可能只有 30%-40%计算单元也经常在等待数据加载整体吞吐远没到瓶颈。这时候再开一个训练任务不是和现有任务抢资源而是把本来就在空转的算力用起来。很多团队算法工程师的标配是一人一张卡但一张卡上同时挂两三个小模型实验是常态这和在公共集群上排队等整机资源不一样单卡多模型更多是为了“不浪费”。另一种常见场景是超参网格搜索。比如一个中小规模的分类模型要试 learning rate、batch size、dropout 的几组组合每个组合上同样跑几十个 epoch。单独串行跑总时长是排队的还经常出现“最后一组结果跑完才能开始下一组”的等待并行跑在单卡上只要显存安排得当整体墙钟时间能缩短一大截。而且单卡并行的代码路径不需要改分布式逻辑调试起来也比多机多卡简单得多。还要提一点容易被忽略的TensorFlow 2 的显存默认是动态增长的虽然沙盒环境里很多时候我们喜欢把显存限制写死但在同时跑多个模型的场景里动态增长模式反而是一个值得优先考虑的机制。它不会在进程启动时立刻占满全卡而是让每个模型按自己的实际需求一点点申请显存相互之间“按需打补丁”。当然前提是你不去手动关闭这个特性以及总需求不要超过物理显存上限。这个机制的取舍后面小节里细说。2. 单进程多模型是首选方案绕开多进程的调度噩梦我在起初试的方案是开两个 Python 进程每个进程里用CUDA_VISIBLE_DEVICES0指向同一张卡各自调model.fit()。想法很简单进程之间天然隔离互不影响。但实际操作中发现TensorFlow 的多进程并行在单卡上比预想的麻烦很多。首先是显存分配。两个进程同时启动时如果其中一个进程把显存吃满另一个进程拿到 0 字节就会直接报错。换成tf.config.experimental.set_memory_growth之后又容易出现两个进程在启动阶段反复试探最终把显存挤爆的尴尬局面。真正的问题是多进程之间没有一个中央协调者去说服谁先谁后TF 只能看到自己进程内的显存申请看不到“隔壁老王的模型也在用同一张卡”。其次多进程会引入很多工程化麻烦进程间要同步进度、每个进程要单独写日志、训练中断后要分别恢复检查点还要处理不必要的模型动态库重复加载。两个进程都往同一个日志目录写东西时文件命名冲突是大概率事件。对我来说这是典型的“用复杂度换隔离”小项目不值得。所以我把方案收敛到了单进程内定义多个模型、交替或分时训练。这样的好处是整个进程的生命周期只初始化一次 CUDA 上下文显存分配统一由 TF 的内存管理器统筹多个模型的train_step可以共用同一个数据管道避免重复加载数据带来的 CPU 开销调试时可以同时看到所有模型的 loss/accuracy 曲线不需要去合并多个日志文件中途想暂停某个模型、调整另一个模型的 batch size直接在主循环里加条件判断即可。这个方案本质上没有牺牲“同时训练”的语义。因为单卡上的 GPU 计算本来就是串行的不存在真正的并行执行tf.function和底层的 CUDA kernel 也只是按照调用顺序排队执行。我们要管理的是“在同一个进程里让多个模型的显存共存并且让训练循环有条不紊地推进”这才是这个题目真正要解决的事情。3. 内存分配策略按需增长、硬上限和预留 buffer 的区别在 TF2 中控制 GPU 显存有三种基本方案默认的按需增长set_memory_growth、直接设置硬上限set_memory_limit、以及把 GPU 划成多个虚拟设备。好多教程只教你gpu_options.allow_growth True这一种但多模型并存场景里三种方案各有适用面。3.1 动态增长模式的优缺点import tensorflow as tf gpus tf.config.experimental.list_physical_devices(GPU) if gpus: for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True)set_memory_growth(gpu, True)会让 TF 在开头尽量少占显存之后按需扩张。这个模式对单模型非常友好对多模型则要小心优点是不需要人为估算每个模型的占用TF 会根据算子实际请求去取显存缺点是一旦某个模型占用了大量显存后面想启动的模型可能就无显存可用或者触发 OOM。更关键的是多个模型在训练过程中的显存需求并不稳定有些层在 forward 时会临时展开很大的中间张量比如 Transformer 的多头注意力有些优化器需要额外的动量 buffer。如果三个模型同时在峰值就可能出现“挤掉”最后一个模型的情况。所以动态增长模式适合模型间峰值错峰明显的场景不适合所有模型同时前向反向的场景。3.2 设置硬上限给每个模型划好“地盘”更可控如果希望每个模型的显存池是隔离的最直接的方式是给进程设置一个显存上限比如 24 GB 的卡一个模型最大只能用到 12 GBgpus tf.config.experimental.list_physical_devices(GPU) if gpus: tf.config.experimental.set_memory_limit(gpus[0], 12 * 1024)这样这个进程无论如何不会超过 12 GB。但要注意set_memory_limit是进程级的它不能在一个进程内分别为两个模型单独限制显存。换句话说如果我一个进程里同时跑两个模型只能设置一个总上限然后靠模型内部的 batch size 和显存清理去控制实际占用。这在实际操作中并不可靠。3.3 虚拟设备划分一个 GPU 拆成两个“假卡”对于“我要模型 A 占用不超过 10 GB模型 B 占用不超过 8 GB”的需求TF2 提供了一种虚拟设备方案physical_gpus tf.config.experimental.list_physical_devices(GPU) if physical_gpus: virtual_config [ tf.config.experimental.VirtualDeviceConfiguration(memory_limit10 * 1024), tf.config.experimental.VirtualDeviceConfiguration(memory_limit8 * 1024), ] tf.config.experimental.set_virtual_device_configuration(physical_gpus[0], virtual_config)设置之后物理的 GPU:0 会被拆成 GPU:0 和 GPU:1 两个虚拟设备。模型 A 在with tf.device(/GPU:0)下定义模型 B 在with tf.device(/GPU:1)下定义它们各自只能使用对应的显存上限。这个方案的优点是显存隔离非常清晰缺点是灵活性差虚拟设备的大小是固定的如果模型 A 的 batch size 临时增大哪怕 GPU:0 的 10 GB 里只用了一半也没有办法把剩余显存借给模型 B。另外虚拟设备机制在部分版本的 TF 里对某些算子支持不够完善自定义 op 或第三方库可能出现兼容性问题。综合来看我最常采用的还是动态增长 主循环内人工控制交替训练的组合不设置硬上限两个模型各自按需申请训练循环里通过调整 batch size 和保留一定余量来防止共峰。只有当模型显存需求差异悬殊、且确实需要严格隔离时才用虚拟设备。4. 核心训练循环实现多个模型在同一进程内的调度方式把多个模型放进同一个进程后训练循环本身并不复杂。关键是处理好三个问题数据供给、梯度更新、以及如何保证两个模型的参数不会互相干扰。下面我从极简的 MNIST 两模型训练代码说起再扩展到更实际的场景。4.1 两个模型的最小并行训练框架import tensorflow as tf import numpy as np # 准备数据 (x_train, y_train), (x_test, y_test) tf.keras.datasets.mnist.load_data() x_train x_train[..., None].astype(np.float32) / 255.0 y_train tf.keras.utils.to_categorical(y_train, 10) train_ds tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(256).prefetch(1) # 定义两个模型 def make_model(): model tf.keras.Sequential([ tf.keras.layers.Flatten(input_shape(28, 28, 1)), tf.keras.layers.Dense(128, activationrelu), tf.keras.layers.Dense(10, activationsoftmax) ]) return model model_a make_model() model_b make_model() opt_a tf.keras.optimizers.Adam(1e-3) opt_b tf.keras.optimizers.Adam(1e-3) loss_fn tf.keras.losses.CategoricalCrossentropy() tf.function def train_step(model, optimizer, x, y): with tf.GradientTape() as tape: pred model(x, trainingTrue) loss loss_fn(y, pred) grads tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) return loss这里最关键的是tf.function的装饰。没有它每一步都会重新执行 Python 级别的运算GPU 利用率会低到让人怀疑人生。加了之后TensorFlow 会把训练步骤编译成图两个模型虽然在同一个进程里但实际执行交给 CUDA 的 kernel 队列训练效率接近串行训练。主训练循环可以这样写epochs 5 for epoch in range(epochs): for step, (x, y) in enumerate(train_ds): # 交替训练每个 step 同时更新两个模型 loss_a train_step(model_a, opt_a, x, y) loss_b train_step(model_b, opt_b, x, y) if step % 100 0: print(fepoch {epoch} step {step}, loss_a{loss_a:.4f}, loss_b{loss_b:.4f})这个“交替训练”的写法很粗糙但它能说明单进程多模型的本质两个模型在同一个 step 上分别前向、反向、更新参数互不干扰。显存占用是两者的叠加计算时间是两者的累加。这也是为什么单卡并行训练的实际提速更多来自“把数据喂得更满”和“减少资源空档”而不是并行计算本身。4.2 进阶按轮次分配训练时间而不是按 step 交替上面这种每个 step 都同时训练两个模型的写法显存峰值必然是两者之和。如果两个模型都比较大很容易把卡打爆。我更常用的是“时分复用”的方式一个 epoch 里先跑模型 A 一段时间再切换到模型 B。这样在同一时刻只有一个人在用 GPU显存峰值反而降下来了。for epoch in range(epochs): # 模型 A 训练一个 epoch for x, y in train_ds: loss_a train_step(model_a, opt_a, x, y) # 切换显存其实不需要切换模型常驻 for x, y in train_ds: loss_b train_step(model_b, opt_b, x, y)但这里有个容易被忽略的问题如果两个模型的参数量都很大比如两个 Transformer即使不训练 AA 的权重也一直占着显存。所以时分复用只能降低“中间激活值”的峰值不能减少参数本身的常驻显存。如果要彻底释放 A 的显存只能del model_a但那就意味着 A 不训练了和我们的目标不符。一个实用折中是把 batch size 调小一点让两者同时训练也撑得住。大多数时候我不精确计算每个张量的字节数直接用一个经验法则显存占用和 batch size 近似成正比。卡是 16GB跑模型 A 单模型时 batch size 64 占 8GB那我给 A 设 batch size 32给 B 设 batch size 32两者加一起大概率在 10GB 以内留下余量给 CUDA context 和碎片。4.3 数据管道要与多个模型解耦多个模型共用同一个tf.data.Dataset时很多新手会把iterator串到一起。比如it iter(train_ds) for x, y in it: loss_a train_step(model_a, ...) loss_b train_step(model_b, ..., x, y)这会导致两个模型的训练进度完全绑定。如果想让 A 一个 epoch 走完整个数据集B 只走一半就需要两个独立的迭代器或重复构建 Datasetit_a iter(train_ds) it_b iter(train_ds)不过那样的话tf.data会为每个迭代器单独维护状态CPU 上可能同时跑两个 prefetch 线程。如果 CPU 核数不多反而会拖慢整体速度。我的习惯是如果两个模型要训练的数据集差别不大就让它们共享同一个迭代器但把目标设计成每步都更新两者如果数据集差异很大比如做迁移学习的不同域才创建两个独立迭代器。另外prefetch的缓冲数量不要设太大。比如prefetch(2)表示 CPU 在 GPU 还在算当前 batch 的时候提前准备接下来两个 batch。两个模型交替吃同一个 prefetch设置过大容易出现 CPU 占用过高或内存溢出。2 到 4 通常已经够了真不够再加。5. 实测显存占用、训练吞吐和几个你没料到的坑把上面的方法用在一张 24GB 的卡上实测了一个图像分类模型和一个文本分类模型的混训。图像模型大概能占 8GB文本模型 6GB两个模型同时训练时显存峰值大概 15GB还有接近 9GB 的冗余。让我意外的是真正卡脖子的不是显存而是数据加载和 CPU 预处理。5.1 数据集加载瓶颈GPU 一直在等数据分别独立跑两个模型时每个模型的数据加载各有自己的 CPU 开销。共享进程后两个模型的map操作会轮流占用 CPU如果都用纯 Python 的预处理函数那么prefetch队列经常填不满。表现为 GPU 利用率只有 60% 左右训练速度比两个模型分别单独跑还慢。解决方法是把预处理尽量向量化成tf.data的原生操作避免大量tf.py_function。比如文本模型要对句子做 tokenize用tf.py_function调用分词库会很慢如果能用tf.strings.split或者预先分词后存成 TFRecord训练速度能明显改善。图像模型的resize、normalize也尽量用tf.image里的原生算子不要自己写 Python 循环。5.2 共峰 OOM 的复现与处理一个比较隐蔽的坑是两个模型的 batch size 在单独调优时都很合适但放在一起后偶尔会在某一步同时产生大中间张量导致 OOM。TF 有时会报错说显存不足但实际可能是 CUDA context 内部碎片化。这个时候把两个模型的 batch size 同时缩小 10%-20%基本能解决。如果缩小后还 OOM那就老老实实用虚拟设备隔开或者把其中一个模型换到 CPU 上跑一小段时间。另外TF2 的显存释放在模型不使用时是比较懒的中间激活缓存不会立刻还给 CUDA。偶尔你会看到显存占用一直在高位其实不代表真的满了。可以用tf.keras.backend.clear_session()强制清理但清理后模型需要重新定义代价有点大只在显存确实不够时用。更温和的方式是tf.config.experimental.enable_memory_usage_logging()去记录内存分配再排查峰值。5.3 断点恢复和日志管理多模型并发最容易翻车的地方多模型同时训练每个模型都要有自己的 checkpoint 目录。我在刚开始写这个方案的时候直接把 A 和 B 的 checkpoint 都放进同一个目录结果 TF 的 checkpoint index 文件互相覆盖导致恢复时模型权重错乱。这算是我踩过最坑的一次。正确的做法是每个模型至少分开两个目录并且带清晰的命名前缀。模型 A 的 checkpoint 放在logs/a/模型 B 的放在logs/b/。如果用tf.train.Checkpoint管理那么每个模型单独声明一个 Checkpoint Managerckpt_a tf.train.Checkpoint(modelmodel_a, optimizeropt_a) manager_a tf.train.CheckpointManager(ckpt_a, directorylogs/a, max_to_keep3) ckpt_b tf.train.Checkpoint(modelmodel_b, optimizeropt_b) manager_b tf.train.CheckpointManager(ckpt_b, directorylogs/b, max_to_keep3)恢复的时候按各自的latest_checkpoint恢复即可。这一点其实和单模型没什么区别但很多人忽略了一个副作用多模型并发训练时如果两个模型都写同一个 TensorBoard 日志文件浏览器里看到的曲线会串行或覆盖。必须使用不同的file_writer目录或者给tf.summary.create_file_writer传入不同的路径。5.4 关于model.fit的局限性如果你打算直接model_a.fit()和model_b.fit()在同一个进程里连续调用很可能发现第二个fit会先把第一个模型的计算图清掉一部分具体看 TF 版本导致模型 B 启动时显存占用异常。TensorFlow 2 的fit封装度高适合单模型快速原型但要精细控制在两块显存块之间切换还是必须自己写训练循环。我自己写训练循环的另一个原因是fit的验证逻辑、early stopping、tensorboard 回调都是对整个进程的全局性操作两个模型共用同一套回调很容易互相干扰。自定义循环虽然代码量多了一点但可控性高得多尤其在多模型并发时你能准确知道每一步在做哪个模型的训练。6. 怎么判断单卡多模型方案值不值边界和替代路径讲了这么多最后说句实在话单卡多模型并不是什么银弹它更适合“小模型叠加”的场景。如果你手头是一个 70B 大模型显存本身就不够再谈同时训练几个模型就没多大意义了。那样的情况下应该先考虑模型并行、流水线并行或者直接换多卡集群。单卡多模型比较适用的边界条件是模型在单卡上能够以“小 batch”完成训练显存占用不超过卡的 40%数据量不大、训练轮数不需要特别长算法侧需要快速并行验证多个模型变体比如对比不同注意力头数、不同 embedding 维度没有条件拿到多卡资源但想通过“填满一张卡”来缩短实验队列。如果超过了这个边界我建议优先考虑两个替代路径一是多进程 多个虚拟设备。我不久前试过在两个进程里分别指定不同的虚拟设备看起来是可行的但需要确保虚拟设备划分必须在一个进程里设置另一个进程要对同一块物理 GPU 做同样的虚拟设备设置。实际做的时候比较绕毕竟CUDA_VISIBLE_DEVICES加上虚拟设备的组合容易让人头晕不是特别必要就别上了。二是使用队列式任务调度器。比如外部写一个 Python 脚本让模型 A 和模型 B 交替从同一个队列取任务把每个模型封装成独立的进程但共享同一张卡。这和前面所说的多进程方案本质接近但有了调度器的约束比裸多进程稳定不少。我自己没有大规模采用这种方案因为单进程多模型在代码上已经够简单直接。另外如果你只是想让模型 A 跑得快一点同时让模型 B 顺便验证一下梯度计算有没有问题那完全可以让 B 占用的显存稍微低一点给它一个小模型比如用常量 batch size 8甚至只用 CPU 计算。这种“主次分明”的策略也能避免不必要的共峰风险。我现在实际跑训练任务时常态是一张 24GB 卡上同时挂一个 12GB 左右的主模型和一个小验证模型。这个验证模型主要用于对比 loss 曲线、检查数据增强的效果训练速度要求不高。用单进程自定义循环两个模型各自输出日志到不同目录跑起来相当稳。如果你也想在单卡上同时塞几个小实验建议第一个版本从“两个小模型交替训练”开始熟悉了数据管道和 checkpoints 的组织方式之后再逐步提高并发度。最后补充一个小技巧在训练循环开始前先用一个很小的 batch 预热每个模型一次把 CUDA kernel 和显存分配提前触发可以避免正式训练时前几十步的抖动。
在训练资源比较紧张的环境里单张 GPU 往往要同时扛好几个实验。我常遇到的情况是白天要跑一个基线模型晚上想再插一个对比实验或者某个大模型把显存占得七七八八剩下的一小块显存既放不下完整 batch又不想让 GPU 完全闲着。TensorFlow 2 在单卡上同时训练多个模型听起来有点“挤牙膏”的意思但真做起来有讲究不是简单把两个 model.fit() 叠在一起就完事。这篇文章就聊我实际试下来比较稳的做法、踩过的坑以及背后的资源分配逻辑给准备在一张卡上同时跑多个模型的同学做个参考。1. 为什么非要在单卡上同时训练资源空转才是最贵的成本先说个很现实的问题一张训练卡如果只跑一个小模型显存利用率可能只有 30%-40%计算单元也经常在等待数据加载整体吞吐远没到瓶颈。这时候再开一个训练任务不是和现有任务抢资源而是把本来就在空转的算力用起来。很多团队算法工程师的标配是一人一张卡但一张卡上同时挂两三个小模型实验是常态这和在公共集群上排队等整机资源不一样单卡多模型更多是为了“不浪费”。另一种常见场景是超参网格搜索。比如一个中小规模的分类模型要试 learning rate、batch size、dropout 的几组组合每个组合上同样跑几十个 epoch。单独串行跑总时长是排队的还经常出现“最后一组结果跑完才能开始下一组”的等待并行跑在单卡上只要显存安排得当整体墙钟时间能缩短一大截。而且单卡并行的代码路径不需要改分布式逻辑调试起来也比多机多卡简单得多。还要提一点容易被忽略的TensorFlow 2 的显存默认是动态增长的虽然沙盒环境里很多时候我们喜欢把显存限制写死但在同时跑多个模型的场景里动态增长模式反而是一个值得优先考虑的机制。它不会在进程启动时立刻占满全卡而是让每个模型按自己的实际需求一点点申请显存相互之间“按需打补丁”。当然前提是你不去手动关闭这个特性以及总需求不要超过物理显存上限。这个机制的取舍后面小节里细说。2. 单进程多模型是首选方案绕开多进程的调度噩梦我在起初试的方案是开两个 Python 进程每个进程里用CUDA_VISIBLE_DEVICES0指向同一张卡各自调model.fit()。想法很简单进程之间天然隔离互不影响。但实际操作中发现TensorFlow 的多进程并行在单卡上比预想的麻烦很多。首先是显存分配。两个进程同时启动时如果其中一个进程把显存吃满另一个进程拿到 0 字节就会直接报错。换成tf.config.experimental.set_memory_growth之后又容易出现两个进程在启动阶段反复试探最终把显存挤爆的尴尬局面。真正的问题是多进程之间没有一个中央协调者去说服谁先谁后TF 只能看到自己进程内的显存申请看不到“隔壁老王的模型也在用同一张卡”。其次多进程会引入很多工程化麻烦进程间要同步进度、每个进程要单独写日志、训练中断后要分别恢复检查点还要处理不必要的模型动态库重复加载。两个进程都往同一个日志目录写东西时文件命名冲突是大概率事件。对我来说这是典型的“用复杂度换隔离”小项目不值得。所以我把方案收敛到了单进程内定义多个模型、交替或分时训练。这样的好处是整个进程的生命周期只初始化一次 CUDA 上下文显存分配统一由 TF 的内存管理器统筹多个模型的train_step可以共用同一个数据管道避免重复加载数据带来的 CPU 开销调试时可以同时看到所有模型的 loss/accuracy 曲线不需要去合并多个日志文件中途想暂停某个模型、调整另一个模型的 batch size直接在主循环里加条件判断即可。这个方案本质上没有牺牲“同时训练”的语义。因为单卡上的 GPU 计算本来就是串行的不存在真正的并行执行tf.function和底层的 CUDA kernel 也只是按照调用顺序排队执行。我们要管理的是“在同一个进程里让多个模型的显存共存并且让训练循环有条不紊地推进”这才是这个题目真正要解决的事情。3. 内存分配策略按需增长、硬上限和预留 buffer 的区别在 TF2 中控制 GPU 显存有三种基本方案默认的按需增长set_memory_growth、直接设置硬上限set_memory_limit、以及把 GPU 划成多个虚拟设备。好多教程只教你gpu_options.allow_growth True这一种但多模型并存场景里三种方案各有适用面。3.1 动态增长模式的优缺点import tensorflow as tf gpus tf.config.experimental.list_physical_devices(GPU) if gpus: for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True)set_memory_growth(gpu, True)会让 TF 在开头尽量少占显存之后按需扩张。这个模式对单模型非常友好对多模型则要小心优点是不需要人为估算每个模型的占用TF 会根据算子实际请求去取显存缺点是一旦某个模型占用了大量显存后面想启动的模型可能就无显存可用或者触发 OOM。更关键的是多个模型在训练过程中的显存需求并不稳定有些层在 forward 时会临时展开很大的中间张量比如 Transformer 的多头注意力有些优化器需要额外的动量 buffer。如果三个模型同时在峰值就可能出现“挤掉”最后一个模型的情况。所以动态增长模式适合模型间峰值错峰明显的场景不适合所有模型同时前向反向的场景。3.2 设置硬上限给每个模型划好“地盘”更可控如果希望每个模型的显存池是隔离的最直接的方式是给进程设置一个显存上限比如 24 GB 的卡一个模型最大只能用到 12 GBgpus tf.config.experimental.list_physical_devices(GPU) if gpus: tf.config.experimental.set_memory_limit(gpus[0], 12 * 1024)这样这个进程无论如何不会超过 12 GB。但要注意set_memory_limit是进程级的它不能在一个进程内分别为两个模型单独限制显存。换句话说如果我一个进程里同时跑两个模型只能设置一个总上限然后靠模型内部的 batch size 和显存清理去控制实际占用。这在实际操作中并不可靠。3.3 虚拟设备划分一个 GPU 拆成两个“假卡”对于“我要模型 A 占用不超过 10 GB模型 B 占用不超过 8 GB”的需求TF2 提供了一种虚拟设备方案physical_gpus tf.config.experimental.list_physical_devices(GPU) if physical_gpus: virtual_config [ tf.config.experimental.VirtualDeviceConfiguration(memory_limit10 * 1024), tf.config.experimental.VirtualDeviceConfiguration(memory_limit8 * 1024), ] tf.config.experimental.set_virtual_device_configuration(physical_gpus[0], virtual_config)设置之后物理的 GPU:0 会被拆成 GPU:0 和 GPU:1 两个虚拟设备。模型 A 在with tf.device(/GPU:0)下定义模型 B 在with tf.device(/GPU:1)下定义它们各自只能使用对应的显存上限。这个方案的优点是显存隔离非常清晰缺点是灵活性差虚拟设备的大小是固定的如果模型 A 的 batch size 临时增大哪怕 GPU:0 的 10 GB 里只用了一半也没有办法把剩余显存借给模型 B。另外虚拟设备机制在部分版本的 TF 里对某些算子支持不够完善自定义 op 或第三方库可能出现兼容性问题。综合来看我最常采用的还是动态增长 主循环内人工控制交替训练的组合不设置硬上限两个模型各自按需申请训练循环里通过调整 batch size 和保留一定余量来防止共峰。只有当模型显存需求差异悬殊、且确实需要严格隔离时才用虚拟设备。4. 核心训练循环实现多个模型在同一进程内的调度方式把多个模型放进同一个进程后训练循环本身并不复杂。关键是处理好三个问题数据供给、梯度更新、以及如何保证两个模型的参数不会互相干扰。下面我从极简的 MNIST 两模型训练代码说起再扩展到更实际的场景。4.1 两个模型的最小并行训练框架import tensorflow as tf import numpy as np # 准备数据 (x_train, y_train), (x_test, y_test) tf.keras.datasets.mnist.load_data() x_train x_train[..., None].astype(np.float32) / 255.0 y_train tf.keras.utils.to_categorical(y_train, 10) train_ds tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(256).prefetch(1) # 定义两个模型 def make_model(): model tf.keras.Sequential([ tf.keras.layers.Flatten(input_shape(28, 28, 1)), tf.keras.layers.Dense(128, activationrelu), tf.keras.layers.Dense(10, activationsoftmax) ]) return model model_a make_model() model_b make_model() opt_a tf.keras.optimizers.Adam(1e-3) opt_b tf.keras.optimizers.Adam(1e-3) loss_fn tf.keras.losses.CategoricalCrossentropy() tf.function def train_step(model, optimizer, x, y): with tf.GradientTape() as tape: pred model(x, trainingTrue) loss loss_fn(y, pred) grads tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) return loss这里最关键的是tf.function的装饰。没有它每一步都会重新执行 Python 级别的运算GPU 利用率会低到让人怀疑人生。加了之后TensorFlow 会把训练步骤编译成图两个模型虽然在同一个进程里但实际执行交给 CUDA 的 kernel 队列训练效率接近串行训练。主训练循环可以这样写epochs 5 for epoch in range(epochs): for step, (x, y) in enumerate(train_ds): # 交替训练每个 step 同时更新两个模型 loss_a train_step(model_a, opt_a, x, y) loss_b train_step(model_b, opt_b, x, y) if step % 100 0: print(fepoch {epoch} step {step}, loss_a{loss_a:.4f}, loss_b{loss_b:.4f})这个“交替训练”的写法很粗糙但它能说明单进程多模型的本质两个模型在同一个 step 上分别前向、反向、更新参数互不干扰。显存占用是两者的叠加计算时间是两者的累加。这也是为什么单卡并行训练的实际提速更多来自“把数据喂得更满”和“减少资源空档”而不是并行计算本身。4.2 进阶按轮次分配训练时间而不是按 step 交替上面这种每个 step 都同时训练两个模型的写法显存峰值必然是两者之和。如果两个模型都比较大很容易把卡打爆。我更常用的是“时分复用”的方式一个 epoch 里先跑模型 A 一段时间再切换到模型 B。这样在同一时刻只有一个人在用 GPU显存峰值反而降下来了。for epoch in range(epochs): # 模型 A 训练一个 epoch for x, y in train_ds: loss_a train_step(model_a, opt_a, x, y) # 切换显存其实不需要切换模型常驻 for x, y in train_ds: loss_b train_step(model_b, opt_b, x, y)但这里有个容易被忽略的问题如果两个模型的参数量都很大比如两个 Transformer即使不训练 AA 的权重也一直占着显存。所以时分复用只能降低“中间激活值”的峰值不能减少参数本身的常驻显存。如果要彻底释放 A 的显存只能del model_a但那就意味着 A 不训练了和我们的目标不符。一个实用折中是把 batch size 调小一点让两者同时训练也撑得住。大多数时候我不精确计算每个张量的字节数直接用一个经验法则显存占用和 batch size 近似成正比。卡是 16GB跑模型 A 单模型时 batch size 64 占 8GB那我给 A 设 batch size 32给 B 设 batch size 32两者加一起大概率在 10GB 以内留下余量给 CUDA context 和碎片。4.3 数据管道要与多个模型解耦多个模型共用同一个tf.data.Dataset时很多新手会把iterator串到一起。比如it iter(train_ds) for x, y in it: loss_a train_step(model_a, ...) loss_b train_step(model_b, ..., x, y)这会导致两个模型的训练进度完全绑定。如果想让 A 一个 epoch 走完整个数据集B 只走一半就需要两个独立的迭代器或重复构建 Datasetit_a iter(train_ds) it_b iter(train_ds)不过那样的话tf.data会为每个迭代器单独维护状态CPU 上可能同时跑两个 prefetch 线程。如果 CPU 核数不多反而会拖慢整体速度。我的习惯是如果两个模型要训练的数据集差别不大就让它们共享同一个迭代器但把目标设计成每步都更新两者如果数据集差异很大比如做迁移学习的不同域才创建两个独立迭代器。另外prefetch的缓冲数量不要设太大。比如prefetch(2)表示 CPU 在 GPU 还在算当前 batch 的时候提前准备接下来两个 batch。两个模型交替吃同一个 prefetch设置过大容易出现 CPU 占用过高或内存溢出。2 到 4 通常已经够了真不够再加。5. 实测显存占用、训练吞吐和几个你没料到的坑把上面的方法用在一张 24GB 的卡上实测了一个图像分类模型和一个文本分类模型的混训。图像模型大概能占 8GB文本模型 6GB两个模型同时训练时显存峰值大概 15GB还有接近 9GB 的冗余。让我意外的是真正卡脖子的不是显存而是数据加载和 CPU 预处理。5.1 数据集加载瓶颈GPU 一直在等数据分别独立跑两个模型时每个模型的数据加载各有自己的 CPU 开销。共享进程后两个模型的map操作会轮流占用 CPU如果都用纯 Python 的预处理函数那么prefetch队列经常填不满。表现为 GPU 利用率只有 60% 左右训练速度比两个模型分别单独跑还慢。解决方法是把预处理尽量向量化成tf.data的原生操作避免大量tf.py_function。比如文本模型要对句子做 tokenize用tf.py_function调用分词库会很慢如果能用tf.strings.split或者预先分词后存成 TFRecord训练速度能明显改善。图像模型的resize、normalize也尽量用tf.image里的原生算子不要自己写 Python 循环。5.2 共峰 OOM 的复现与处理一个比较隐蔽的坑是两个模型的 batch size 在单独调优时都很合适但放在一起后偶尔会在某一步同时产生大中间张量导致 OOM。TF 有时会报错说显存不足但实际可能是 CUDA context 内部碎片化。这个时候把两个模型的 batch size 同时缩小 10%-20%基本能解决。如果缩小后还 OOM那就老老实实用虚拟设备隔开或者把其中一个模型换到 CPU 上跑一小段时间。另外TF2 的显存释放在模型不使用时是比较懒的中间激活缓存不会立刻还给 CUDA。偶尔你会看到显存占用一直在高位其实不代表真的满了。可以用tf.keras.backend.clear_session()强制清理但清理后模型需要重新定义代价有点大只在显存确实不够时用。更温和的方式是tf.config.experimental.enable_memory_usage_logging()去记录内存分配再排查峰值。5.3 断点恢复和日志管理多模型并发最容易翻车的地方多模型同时训练每个模型都要有自己的 checkpoint 目录。我在刚开始写这个方案的时候直接把 A 和 B 的 checkpoint 都放进同一个目录结果 TF 的 checkpoint index 文件互相覆盖导致恢复时模型权重错乱。这算是我踩过最坑的一次。正确的做法是每个模型至少分开两个目录并且带清晰的命名前缀。模型 A 的 checkpoint 放在logs/a/模型 B 的放在logs/b/。如果用tf.train.Checkpoint管理那么每个模型单独声明一个 Checkpoint Managerckpt_a tf.train.Checkpoint(modelmodel_a, optimizeropt_a) manager_a tf.train.CheckpointManager(ckpt_a, directorylogs/a, max_to_keep3) ckpt_b tf.train.Checkpoint(modelmodel_b, optimizeropt_b) manager_b tf.train.CheckpointManager(ckpt_b, directorylogs/b, max_to_keep3)恢复的时候按各自的latest_checkpoint恢复即可。这一点其实和单模型没什么区别但很多人忽略了一个副作用多模型并发训练时如果两个模型都写同一个 TensorBoard 日志文件浏览器里看到的曲线会串行或覆盖。必须使用不同的file_writer目录或者给tf.summary.create_file_writer传入不同的路径。5.4 关于model.fit的局限性如果你打算直接model_a.fit()和model_b.fit()在同一个进程里连续调用很可能发现第二个fit会先把第一个模型的计算图清掉一部分具体看 TF 版本导致模型 B 启动时显存占用异常。TensorFlow 2 的fit封装度高适合单模型快速原型但要精细控制在两块显存块之间切换还是必须自己写训练循环。我自己写训练循环的另一个原因是fit的验证逻辑、early stopping、tensorboard 回调都是对整个进程的全局性操作两个模型共用同一套回调很容易互相干扰。自定义循环虽然代码量多了一点但可控性高得多尤其在多模型并发时你能准确知道每一步在做哪个模型的训练。6. 怎么判断单卡多模型方案值不值边界和替代路径讲了这么多最后说句实在话单卡多模型并不是什么银弹它更适合“小模型叠加”的场景。如果你手头是一个 70B 大模型显存本身就不够再谈同时训练几个模型就没多大意义了。那样的情况下应该先考虑模型并行、流水线并行或者直接换多卡集群。单卡多模型比较适用的边界条件是模型在单卡上能够以“小 batch”完成训练显存占用不超过卡的 40%数据量不大、训练轮数不需要特别长算法侧需要快速并行验证多个模型变体比如对比不同注意力头数、不同 embedding 维度没有条件拿到多卡资源但想通过“填满一张卡”来缩短实验队列。如果超过了这个边界我建议优先考虑两个替代路径一是多进程 多个虚拟设备。我不久前试过在两个进程里分别指定不同的虚拟设备看起来是可行的但需要确保虚拟设备划分必须在一个进程里设置另一个进程要对同一块物理 GPU 做同样的虚拟设备设置。实际做的时候比较绕毕竟CUDA_VISIBLE_DEVICES加上虚拟设备的组合容易让人头晕不是特别必要就别上了。二是使用队列式任务调度器。比如外部写一个 Python 脚本让模型 A 和模型 B 交替从同一个队列取任务把每个模型封装成独立的进程但共享同一张卡。这和前面所说的多进程方案本质接近但有了调度器的约束比裸多进程稳定不少。我自己没有大规模采用这种方案因为单进程多模型在代码上已经够简单直接。另外如果你只是想让模型 A 跑得快一点同时让模型 B 顺便验证一下梯度计算有没有问题那完全可以让 B 占用的显存稍微低一点给它一个小模型比如用常量 batch size 8甚至只用 CPU 计算。这种“主次分明”的策略也能避免不必要的共峰风险。我现在实际跑训练任务时常态是一张 24GB 卡上同时挂一个 12GB 左右的主模型和一个小验证模型。这个验证模型主要用于对比 loss 曲线、检查数据增强的效果训练速度要求不高。用单进程自定义循环两个模型各自输出日志到不同目录跑起来相当稳。如果你也想在单卡上同时塞几个小实验建议第一个版本从“两个小模型交替训练”开始熟悉了数据管道和 checkpoints 的组织方式之后再逐步提高并发度。最后补充一个小技巧在训练循环开始前先用一个很小的 batch 预热每个模型一次把 CUDA kernel 和显存分配提前触发可以避免正式训练时前几十步的抖动。
平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测)
