在训练资源比较紧张的环境里，单张 GPU 往往要同时扛好几个实验。我常遇到的情况是：白天要跑一个基线模型，晚上想再插一个对比实验；或者某个大模型把显存占得七七八八，剩下的一小块显存既放不下完整 batch，又不想让 GPU 完全闲着。TensorFlow 2 在单卡上同时训练多个模型，听起来有点“挤牙膏”的意思，但真做起来有讲究，不是简单把两个 model.fit() 叠在一起就完事。这篇文章就聊我实际试下来比较稳的做法、踩过的坑，以及背后的资源分配逻辑，给准备在一张卡上同时跑多个模型的同学做个参考。
1. 为什么非要在单卡上同时训练：资源空转才是最贵的成本
先说个很现实的问题：一张训练卡如果只跑一个小模型，显存利用率可能只有 30%-40%，计算单元也经常在等待数据加载，整体吞吐远没到瓶颈。这时候再开一个训练任务，不是和现有任务抢资源，而是把本来就在空转的算力用起来。很多团队算法工程师的标配是一人一张卡，但一张卡上同时挂两三个小模型实验是常态，这和在公共集群上排队等整机资源不一样，单卡多模型更多是为了“不浪费”。
另一种常见场景是超参网格搜索。比如一个中小规模的分类模型，要试 learning rate、batch size、dropout 的几组组合，每个组合上同样跑几十个 epoch。单独串行跑，总时长是排队的，还经常出现“最后一组结果跑完才能开始下一组”的等待；并行跑在单卡上，只要显存安排得当，整体墙钟时间能缩短一大截。而且单卡并行的代码路径不需要改分布式逻辑，调试起来也比多机多卡简单得多。
还要提一点容易被忽略的：TensorFlow 2 的显存默认是动态增长的，虽然沙盒环境里很多时候我们喜欢把显存限制写死，但在同时跑多个模型的场景里，动态增长模式反而是一个值得优先考虑的机制。它不会在进程启动时立刻占满全卡，而是让每个模型按自己的实际需求一点点申请显存，相互之间“按需打补丁”。当然，前提是你不去手动关闭这个特性，以及总需求不要超过物理显存上限。这个机制的取舍后面小节里细说。
2. 单进程多模型是首选方案：绕开多进程的调度噩梦
我在起初试的方案是开两个 Python 进程，每个进程里用CUDA_VISIBLE_DEVICES=0指向同一张卡，各自调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 GB：
gpus = 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_limit=10 * 1024), tf.config.experimental.VirtualDeviceConfiguration(memory_limit=8 * 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, activation='relu'), tf.keras.layers.Dense(10, activation='softmax') ]) 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, training=True) 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(f"epoch {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），即使不训练 A，A 的权重也一直占着显存。所以时分复用只能降低“中间激活值”的峰值，不能减少参数本身的常驻显存。如果要彻底释放 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 只走一半，就需要两个独立的迭代器或重复构建 Dataset：
it_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 Manager：
ckpt_a = tf.train.Checkpoint(model=model_a, optimizer=opt_a) manager_a = tf.train.CheckpointManager(ckpt_a, directory="logs/a", max_to_keep=3) ckpt_b = tf.train.Checkpoint(model=model_b, optimizer=opt_b) manager_b = tf.train.CheckpointManager(ckpt_b, directory="logs/b", max_to_keep=3)恢复的时候，按各自的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 和显存分配提前触发，可以避免正式训练时前几十步的抖动。
