一个多任务模型同时跑分类、分割、深度回归三条线，单卡还能勉强凑合，一旦上到四卡八卡，最先崩的往往不是显存，而是训练节奏——某个任务的 loss 忽然反弹，某个 rank 直接 hang 住，或者吞吐量比单任务还低。PyTorch 的多任务学习和并行处理这两件事，单独拿出来都不算难，难的是它们叠在一起：多任务带来了不平衡的 loss 量纲、稀疏的任务分支和梯度的相互拉扯，并行处理又把这些复杂性复制到每个进程、每个通信桶里，任何一个细节没对齐，都会被放大成"训练不收敛"。我这两年把一套多任务训练框架从单卡折腾到多机，踩过的坑基本都集中在"任务调度"和"通信同步"这两个交界处。

下面这些内容，适合已经能用 PyTorch 跑通单任务训练、准备把多个任务塞进同一个模型、并且需要靠多卡多进程把训练时间压下来的朋友。我会从整体设计思路讲起，再到网络结构、损失加权、数据并行落地、显存优化，最后把排查过程整理成能直接查的表。文中涉及的模型和数据集都是我搭的模拟项目，思路和参数可以直接迁移。

1. 整体设计与思路拆解：为什么要这么组合

1.1 多任务学习真正的价值点在哪里

先把话说清楚：多任务学习不是"省一个模型"这么简单。它的核心收益在于共享表征带来的正则化效应和标注数据的复用。举个我实际遇到的情况——一套图像处理 Demo，需要同时输出区域分类、前景分割掩码和目标距离估计。如果拆成三个独立模型，等于三份 backbone、三倍的前向开销，而且分类任务标注量大、距离估计标注量小，小任务单独训很容易过拟合。共享 backbone 之后，分类带来的大量梯度给小任务的浅层特征提供了更丰富的监督信号，小任务的泛化误差下来了不少。

但这里有个容易被忽略的前提：任务之间必须存在可共享的底层特征。分类靠的是语义，分割靠的是边界和纹理，距离估计靠的是尺度线索，这三者在浅层是共通的（边缘、颜色、梯度），在深层就开始分化。所以我选择的是"浅层硬共享、深层各自分支"的结构，而不是从第一层就分开。如果任务之间毫无关联，比如把图像分类和音频分类硬凑一起，共享层反而互相拖累，这就是常说的负迁移。

判断要不要上多任务，我的经验是看三条：任务输入是不是同一模态、标注能不能对齐到同一个样本粒度、单任务数据量是否都够。三条里有一条不满足，先做单任务再考虑融合，别一上来就搞多任务，调参成本会把你劝退。

1.2 并行处理的三条路线怎么选

PyTorch 侧能用的并行手段大致分三类，选错了会白折腾很久。

多任务场景下，我基本是"数据并行打底 + 显存优化兜底"的组合。原因很实在：多任务模型通常是共享 backbone 加多个小 head，head 参数量很小，瓶颈在 backbone 的激活值，不在参数本身。所以先用 DDP 把吞吐拉起来，再用混合精度和梯度累积压显存，比一上来就上 FSDP 划算得多。FSDP 的切分和通信开销是有代价的，参数规模没到那个量级，纯属给自己添堵。

真正需要切分模型的信号很明确：单个 rank 上 backbone 的激活值就撑爆显存，且减小 batch 已经影响到 BatchNorm 的统计稳定性。这时候再考虑 FSDP 的SHARDING_STRATEGY.FULL_SHARD，或者对 backbone 做层级的流水线切分。

1.3 组合方案的取舍逻辑

把多任务和并行叠在一起，最需要提前定下来的是任务采样策略和梯度同步粒度这两件事，它们决定了后面所有代码的写法。

任务采样上，我试过两种。一种是"每步混合采样"——每个 batch 里按比例塞进不同任务的样本，一次前向拿到所有任务的输出。好处是通信效率高，一次 all-reduce 同步全部梯度；坏处是不同任务的数据集要能拼到一起，且 batch 内的有效样本数被摊薄。另一种是"按步轮转"——第 n 步只训任务 A，第 n+1 步只训任务 B。好处是每个任务的梯度是干净的，坏处是每步都要同步，通信次数翻倍，而且某些 head 在当步没有梯度，DDP 会报未使用参数的错误。

我最后选的是混合采样，理由是通信开销在多卡下是硬成本，而混合采样天然适配 DDP 的一次性梯度同步。代价是要处理"某任务在当前 batch 里没有样本"的情况，得用 mask 把对应的 loss 屏蔽掉，这个后面会细讲。

梯度同步粒度上，我坚持一个原则：所有任务的 loss 必须先在本 rank 内聚合成一个标量，再做 backward。不要对每个任务分别 backward，那样会触发多次 all-reduce，通信量随任务数线性增长。聚合的方式就是加权求和，权重怎么定，是第 2 章的重点。

2. 多任务学习的核心细节与结构实现

2.1 硬共享、软共享与专家混合的取舍

三种主流范式我都在模拟项目里跑过，结论是：硬共享是默认答案，软共享是备选，专家混合是特定场景的解药。

硬共享就是共享底层、分叉高层。它的优势是参数量最省、推理最快、实现最简单。缺点是任务冲突不可控，一旦某个任务的梯度主导了共享层，其他任务就只能被动接受。我一般会在硬共享的基础上做两件事来缓解：一是给每个 head 加一个独立的 BN 或者 LayerNorm，让各任务有自己的归一化统计量；二是让 head 的容量足够，别用一个线性层草草了事，两到三层的 MLP 往往能显著缓解 head 之间的干扰。

软共享是每个任务各有一套完整网络，在 loss 里加一个约束项，让不同网络的中间特征保持接近，比如对特征图做 L2 距离约束，或者对注意力做一致性约束。它的好处是任务间干扰小，坏处是参数量和显存都翻倍，多卡下更吃紧。我一般只在任务差异极大、硬共享明显负迁移时才用。

专家混合（MoE 那一路，比如 MMoE、PLE）的思路是给每个任务配一组专家，通过门控网络加权组合。它适合任务之间既有共性又有明显差异的场景，比如"用户行为预测"里点击率和转化率的关系。代价是门控网络本身需要训练，样本少的时候门控容易塌缩到只用一两个专家，实际退化成硬共享。我在这套图像 Demo 里没上 MoE，因为任务数只有三个，硬共享加独立归一化已经够用。

提示：判断要不要升级范式，别看论文里的指标，看你自己项目里单任务 baseline 和多任务版本的差距。如果多任务版本在某个任务上比单任务差超过 3 个点，才值得考虑换结构。

2.2 损失权重与梯度冲突：最容易翻车的地方

多任务训练里，我见过最多的翻车不是结构设计，而是 loss 量纲没对齐。分类用交叉熵，量级在 0.5 到 3 之间；分割用 Dice 或者 BCE，量级可能是 0.1；距离回归用 L1，量级取决于标签尺度，可能是几十。三个 loss 直接相加，回归任务直接主导整个训练，分类和分割基本学不动。

最朴素的解法是手工调权重，我早期就是这么干的，效果是"调好了这个任务，那个任务又崩了"，非常折磨。后来换成基于同方差不确定性的自动加权，稳定性好了很多。原理不复杂：假设每个任务的观测噪声是高斯分布，噪声越大，这个任务的权重就应该越小，而噪声本身作为一个可学习参数参与训练。公式上每个任务的项变成L_i / (2σ_i²) + log σ_i，实现时用log σ_i²作为参数，取指数就是精度。

这个方法的妙处在于它是自适应的：训练初期回归任务 loss 大，σ 会被推大，权重自动降下来；等回归任务收敛了，σ 回落，权重又升上去。但它不是万能的，如果某个任务本身就是病态的（标签噪声极大），它会把这个任务的权重压到接近零，相当于放弃了这个任务。所以我会给它加一个下限约束，或者对 σ 做 clamp。

另一个必须处理的问题是梯度冲突。多任务共享层收到的梯度是各任务梯度的叠加，如果两个任务的梯度方向相反，就会互相抵消。检测方法很直接：在共享层上分别算每个任务的梯度，两两算余弦相似度，如果长期为负，说明存在冲突。缓解手段我用过两种，一种是 PCGrad 的思路——当两个梯度夹角大于 90 度时，把其中一个投影到另一个的法平面上，去掉冲突分量；另一种是 GradNorm，动态调整权重让各任务的梯度范数保持接近。

PCGrad 的实现不复杂，核心就几行：

def project_conflicting(grads): # grads: list of list[Tensor]，每个任务的梯度列表 n = len(grads) for i in range(n): for j in range(n): if i == j: continue dot = sum((gi * gj).sum() for gi, gj in zip(grads[i], grads[j])) if dot < 0: norm_j = sum((gj ** 2).sum() for gj in grads[j]) + 1e-12 coef = dot / norm_j grads[i] = [gi - coef * gj for gi, gj in zip(grads[i], grads[j])] return grads

这段代码的代价是每个任务都要单独算一次梯度，显存和时间都会涨。我的经验是：只有在检测到明确冲突时才开 PCGrad，而且只对共享 backbone 的最后几层开，全网络开启性价比太低。

2.3 一个可复现的多任务模型骨架

结构上我遵循"共享主干 + 独立归一化 head + 统一输出字典"的写法。输出统一成 dict，后面无论是算 loss 还是做 DDP 封装都清爽。

import torch import torch.nn as nn class SharedBackbone(nn.Module): def __init__(self, in_ch=3, width=32): super().__init__() self.stem = nn.Sequential( nn.Conv2d(in_ch, width, 3, stride=2, padding=1), nn.BatchNorm2d(width), nn.ReLU(inplace=True), nn.Conv2d(width, width * 2, 3, stride=2, padding=1), nn.BatchNorm2d(width * 2), nn.ReLU(inplace=True), nn.Conv2d(width * 2, width * 4, 3, stride=2, padding=1), nn.BatchNorm2d(width * 4), nn.ReLU(inplace=True), ) self.out_ch = width * 4 def forward(self, x): return self.stem(x) class MultiTaskNet(nn.Module): def __init__(self, num_classes=10, head_hidden=128): super().__init__() self.backbone = SharedBackbone() c = self.backbone.out_ch self.head_cls = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(c, head_hidden), nn.LayerNorm(head_hidden), nn.ReLU(inplace=True), nn.Linear(head_hidden, num_classes), ) self.head_seg = nn.Sequential( nn.Conv2d(c, head_hidden, 3, padding=1), nn.BatchNorm2d(head_hidden), nn.ReLU(inplace=True), nn.Conv2d(head_hidden, 1, 1), ) self.head_depth = nn.Sequential( nn.Conv2d(c, head_hidden, 3, padding=1), nn.BatchNorm2d(head_hidden), nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(head_hidden, 1), ) def forward(self, x): feat = self.backbone(x) return { "cls": self.head_cls(feat), "seg": self.head_seg(feat), "depth": self.head_depth(feat), }

这里有几个刻意的设计。分类 head 用 LayerNorm 而不是 BatchNorm，因为它接在全局池化之后，维度是 1x1，BN 在小 batch 下统计量不稳。分割 head 保留空间分辨率，所以用 BN。深度回归 head 用全局池化后接线性层，回归目标是一个标量，不需要空间输出。三段 head 的容量都比"一个线性层"要大一些，这是为了减少任务之间的干扰——head 越浅，共享层的梯度冲突越容易直接传导。

损失封装我用一个模块统管，把自动加权和 mask 逻辑都放进去：

import torch import torch.nn as nn import torch.nn.functional as F class MTLLoss(nn.Module): def __init__(self, task_names, use_uncertainty=True, min_log_var=-4.0, max_log_var=4.0): super().__init__() self.task_names = task_names self.use_uncertainty = use_uncertainty self.min_log_var = min_log_var self.max_log_var = max_log_var if use_uncertainty: self.log_var = nn.ParameterDict({ n: nn.Parameter(torch.zeros(())) for n in task_names }) def forward(self, outputs, targets, masks): raw = {} if masks["cls"].any(): raw["cls"] = F.cross_entropy( outputs["cls"][masks["cls"]], targets["cls"][masks["cls"]] ) if masks["seg"].any(): raw["seg"] = F.binary_cross_entropy_with_logits( outputs["seg"][masks["seg"]].squeeze(1), targets["seg"][masks["seg"]].float(), ) if masks["depth"].any(): raw["depth"] = F.l1_loss( outputs["depth"][masks["depth"]].squeeze(1), targets["depth"][masks["depth"]].float(), ) total = 0.0 detail = {} for n, loss in raw.items(): if self.use_uncertainty: lv = torch.clamp(self.log_var[n], self.min_log_var, self.max_log_var) total = total + torch.exp(-lv) * loss + lv else: total = total + loss detail[n] = loss.detach() return total, detail

masks是每个任务在当批是否有样本的布尔掩码。这个设计在多卡混合采样下是必须的，否则某个 rank 上某任务一个样本都没有，对应的 head 就没有梯度，DDP 默认会报错。这一点后面还会展开。

3. 并行处理的落地实操

3.1 多进程 DataLoader：参数不是随便填的

多卡训练首先卡住很多人的是数据加载。num_workers=0时数据加载和计算串行，GPU 利用率常年个位数；num_workers开太大，内存爆掉或者进程调度开销反而拖慢。

我的经验公式是：num_workers取min(8, CPU物理核数 // (GPU数 * 每GPU进程数))，然后在这个基础上做实测微调。比如一台机器 32 物理核、8 张卡、每卡一个进程，那每个进程分到 4 个 worker，正好。如果开了超线程，按物理核算，别按逻辑核算，否则上下文切换的代价会吃掉收益。

其余几个参数的作用和取值要一起看：

drop_last=True在多卡下不是可选项而是必须项。如果各 rank 的最后一批样本数不同，一个 rank 已经走到下一个 epoch，另一个还在等同步，整个训练就会 hang 住，而且不报错，只表现为卡死。我最早遇到这个问题时排查了整整一天，最后靠打印每个 rank 的 step 计数才发现。

还有一个容易忽略的点：persistent_workers=True必须配合num_workers>0，而且在 epoch 之间数据集状态（比如随机种子）的更新要放到 worker_init_fn 里，否则每个 epoch 采出来的顺序是一样的。

3.2 DDP 的标准落地步骤

DDP 的接入流程我固化成五步，照着走基本不会出结构性问题。

第一步，初始化进程组，绑定当前进程的 GPU：

import os import torch import torch.distributed as dist def setup_dist(): dist.init_process_group(backend="nccl") local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(local_rank) return local_rank, dist.get_rank(), dist.get_world_size() def cleanup_dist(): dist.destroy_process_group()

第二步，构造 dataset 和 sampler。多任务场景下我用的是按任务分别建 sampler，再组合：

from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler def build_loaders(task_datasets, batch_size, world_size, rank): loaders = {} samplers = {} for name, ds in task_datasets.items(): sampler = DistributedSampler( ds, num_replicas=world_size, rank=rank, shuffle=True, drop_last=True ) samplers[name] = sampler loaders[name] = DataLoader( ds, batch_size=batch_size, sampler=sampler, num_workers=4, pin_memory=True, persistent_workers=True, drop_last=True, ) return loaders, samplers

第三步，模型包 DDP。这里的关键参数是find_unused_parameters：

from torch.nn.parallel import DistributedDataParallel as DDP model = MultiTaskNet().cuda() model = DDP( model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True, gradient_as_bucket_view=True, broadcast_buffers=True, )

find_unused_parameters=True在多任务里几乎躲不开。因为当某个 batch 里某个任务没有样本时，对应 head 的参数不参与反向传播，DDP 默认会认为这些参数"未使用"而报错。开启之后 DDP 会额外做一次图遍历，性能损失大概在 5% 到 10%，但换来的是稳定性。如果能把采样策略设计成"每个 rank 每步都有所有任务的样本"，就可以关掉它，我当时对比过，关掉后吞吐提升约 8%。

gradient_as_bucket_view=True是我建议默认打开的，它让梯度的 bucket 直接作为视图，省掉一次显存拷贝，能明显降低显存峰值。broadcast_buffers=True对含 BN 的模型必须开，否则各卡上的 running_mean 和 running_var 会逐渐发散。

第四步，每个 epoch 前调sampler.set_epoch(epoch)。这行代码漏掉的话，所有 epoch 的样本顺序完全一样，分布式随机性等于没有。多任务下每个任务的 sampler 都要调。

第五步，日志和保存只在 rank 0 做：

if rank == 0: torch.save({"model": model.module.state_dict()}, "ckpt.pt")

用model.module拿原始模型，否则 state_dict 的 key 会带module.前缀，后面加载容易出问题。

启动用的命令是torchrun：

torchrun --nproc_per_node=4 --nnodes=1 --node_rank=0 \ --master_addr=127.0.0.1 --master_port=29500 \ train_mtl.py

--nproc_per_node是每台机器的进程数，通常等于 GPU 数。多机的话每台机器都要跑这条命令，node_rank依次递增，master_addr指向 0 号机器。

3.3 显存优化与吞吐调优

多任务模型的显存开销主要来自三个方面：参数、优化器状态、激活值。前两者是固定的，激活值随 batch 线性增长，所以调优的主战场在激活值。

混合精度是第一优先的手段。我用的写法是 PyTorch 2.x 的torch.amp：

scaler = torch.amp.GradScaler("cuda") for step, batch in enumerate(loader): optimizer.zero_grad(set_to_none=True) with torch.amp.autocast("cuda", dtype=torch.float16): outputs = model(batch["image"]) loss, detail = criterion(outputs, batch["targets"], batch["masks"]) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) scaler.step(optimizer) scaler.update()

这里有三个细节值得说。optimizer.zero_grad(set_to_none=True)比默认的置零更快，因为它只是把.grad置为 None，省掉了写零的操作。unscale_之后再做梯度裁剪是必须的，因为 scaler 会把梯度放大，直接裁剪等于裁错了尺度。clip_grad_norm_在多任务下我固定开，原因是自动加权会让不同阶段的 loss 量级剧烈变化，梯度范数跳变是常态，不裁剪容易炸。

如果显存还是不够，第二手段是梯度累积。注意累积时 loss 要除以累积步数，否则等效学习率会放大：

accum_steps = 4 loss = loss / accum_steps scaler.scale(loss).backward() if (step + 1) % accum_steps == 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True)

第三手段是激活检查点，对 backbone 的每个 block 用torch.utils.checkpoint.checkpoint包一层。代价是前向要重算一次，时间换空间，大约多 30% 的时间换 40% 到 50% 的激活显存。我一般只在输入分辨率很高时才开。

batch size 与学习率的对应关系要单独说。DDP 下有效 batch 是单卡 batch 乘以卡数，如果沿用单卡的 lr，学习率偏小，收敛会明显变慢。线性缩放规则是lr = base_lr * (global_batch / base_batch)，base_batch 取 256、base_lr 取 0.1 是常见的组合，但实际用的时候必须加 warmup，否则训练前几百步容易发散。我通常用 500 到 1000 步的线性 warmup，之后接余弦退火。

4. 常见问题与排查实录

4.1 训练不收敛与任务互相拖累的排查路径

多任务最典型的问题是"总 loss 在降，但某个任务的指标不动甚至变差"。我的排查顺序是固定的四步。

第一步，看每个任务的原始 loss 曲线，不看加权后的总 loss。自动加权会把某些任务的权重压到很小，总 loss 看着很漂亮，实际上那个任务已经放弃了。我在 tensorboard 里给每个任务的原始 loss 单独开一条线，这是判断问题的第一手证据。

第二步，跟单任务 baseline 对比。给同一个 backbone 单独训一个分类模型，如果多任务的分类指标比单任务差 5 个点以上，说明存在明显的负迁移，这时候优先怀疑结构和权重，而不是继续调学习率。

第三步，检查梯度余弦相似度。我在共享 backbone 的最后一层挂了个 hook，每 100 步算一次各任务梯度的两两余弦相似度。如果长期低于 -0.2，就是明确冲突，考虑开 PCGrad 或者调整采样比例。

第四步，检查数据采样比例。这一步很多人会忽略。如果分类数据集有 10 万张，分割只有 5000 张，按自然比例混合采样，分类的梯度会压倒性地主导共享层。我的做法是按1/sqrt(N_i)或者N_i^0.3的幂次做温度采样，把大任务的采样概率压下来。实测下来，把 10:1 的数据量差距压到 3:1 左右的采样比例，小任务的指标能提升 4 到 6 个点。

4.2 分布式训练卡死、OOM 与吞吐不达标的处理

DDP 相关的问题我整理成了一张速查表，按现象查就行。

有几个点我想单独强调。卡在第一步不动，八成是端口冲突，尤其在共享机器上，29500 这个默认端口经常被占。我现在的习惯是每次启动前随机挑一个 20000 到 30000 之间的端口。

中途 hang 住最隐蔽。除了drop_last，还有一种情况是代码里有if loss > threshold: break之类的条件，某个 rank 提前跳出循环，其他 rank 还在等 all-reduce。DDP 里的控制流必须保证所有 rank 一致，任何基于 loss 值的分支都要用dist.all_reduce先做全局同步再判断。

吞吐不达标，我一般的判断方法是看 GPU 利用率的波形。如果利用率呈现规律的锯齿（高一段低一段），说明数据加载跟不上，加 worker；如果利用率一直在 60% 到 70% 徘徊不上去，可能是通信开销占比太高，考虑加大 batch 让计算通信比上升，或者检查是不是开了find_unused_parameters带来的图遍历开销。

4.3 几个踩坑后才明白的细节

NCCL 的亲和性设置。多卡机器上如果没有正确绑定 CPU 和 GPU 的 NUMA 关系，跨 NUMA 访问会让通信带宽掉一截。我用numactl或者启动脚本里设OMP_NUM_THREADS，把每个进程绑到对应的 NUMA 节点，实测通信时间能降 10% 左右。这个不是必选项，但在大规模下有感知。

BN 在多任务里的隐性干扰。共享 backbone 里的 BN 层，其 running statistics 是所有任务的数据混合统计出来的。如果任务间的数据分布差异大（比如不同来源的图像亮度差很多），混合统计量对每个任务都不是最优。我的处理是把 backbone 的 BN 改成 GroupNorm，或者按任务做条件 BN。前者改动小、通用性好，后者效果好但要改结构，需要根据任务差异程度决定。

检查点保存的频率与开销。多任务模型参数多，保存一次可能几百 MB。如果每个 epoch 都存，在多机共享存储上会成为瓶颈，甚至阻塞训练。我改成每 N 个 epoch 存一次，且只在 rank 0 上异步保存到本地盘再后台同步，训练本身不受影响。

评估阶段的 sampler。训练用DistributedSampler，评估也必须用，而且要shuffle=False并且保证每个样本只被一个 rank 处理。如果直接用普通 DataLoader，每个 rank 都会跑全量验证集，时间乘以卡数。我见过有人为了省事在 rank 0 上做全量评估，其他 rank 空等，结果是评估阶段 GPU 利用率掉到 5%，几卡的机器评估比单卡还慢。

4.4 参数选择的经验数值汇总

最后把几个关键参数的经验值整理一下，这些是我在几个模拟项目里反复验证过的起点，实际项目里要在此基础上微调。

这套参数不是最优解，但能保证训练跑起来不崩、指标不塌。多任务和并行处理的调优本质上是个系统工程，结构、损失、采样、通信四个环节里任何一个没对齐，表现出来的都是"不收敛"，而排查时最有效的手段永远是把每个任务的原始指标单独打出来看，别被聚合后的总 loss 骗了。

我在最近这个图像处理 Demo 上最后的配置是四卡 DDP、混合精度、不确定性加权、按 0.3 幂次温度采样，三个任务的指标都超过了各自的单任务 baseline，其中数据量最小的深度回归任务提升最明显。这个结果也印证了一件事：多任务的价值不在于省模型，而在于让小任务借到大任务的表征能力，前提是你得把权重和采样这两件事调对。