ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

推荐系统GPU优化实战:Padding与Mask处理不规则序列

推荐系统GPU优化实战:Padding与Mask处理不规则序列 1. 核心矛盾解析GPU 为什么天生和推荐序列“不对付”1.1 不规则序列在推荐场景里到底长什么样做推荐系统的同学应该都有这种体会拿到的用户行为数据几乎不可能干干净净地躺成一张表格。用户A昨天看了 3 个商品用户B过去七天点了 50 个商品用户C是刚注册的新客行为序列压根就是空的。再加上不同内容类型的曝光时长、停留时间、点击位置这些附加信息序列里每个元素本身还不是一个标量而是一个多维特征向量。这种序列在数据层面体现出来就是经典的“锯齿状”结构。用Python的list来存每个batch里各条样本的长度参差不齐如果转成numpy数组再强行堆成矩阵短的样本尾部必须补上一些无效值否则形状就对不齐。到这一步还没什么大问题因为数据处理阶段本来就是逐条做的CPU单线程或者多进程都能扛得住。真正让人头疼的是把数据送进GPU的那一刻。GPU是一个非常“固执”的运算单元它要求所有参与并行的张量都是规则的矩形——形状固定、数据类型统一、内存布局整齐。不规则序列直接塞进去最常见的报错就是Expected tensor for argument #1 input to have the same dimension或者干脆就是在collate_fn阶段抛出一堆形状不匹配的异常。1.2 GPU 的并行模型决定了它吃“方”不吃“圆”要理解为什么GPU这么挑食得从它的硬件架构讲起。GPU的算力核心分成一组一组的Streaming Multiprocessor每个SM内部又以SIMT单指令多线程的方式执行任务。所谓SIMT简单说就是同一时刻一个SM里的所有线程必须执行同一条指令只是操作的数据不同。这种架构下最理想的计算场景是所有线程都在做同一件事数据可以整齐地切成等长的小块每个线程拿一块算完再拼回来。深度学习里最常见的卷积和矩阵乘法之所以能被优化到极致就是因为它们的计算模式非常规整完全契合SIMT的执行方式。而不规则序列的问题在于如果你让一个线程处理一条样本那每个线程要处理的数据长度不一样。短样本的线程算完了得干等着长样本的线程SM里大量的执行单元处于空闲状态。更麻烦的是不同长度导致的内存访问地址是跳跃的没法做合并访存显存带宽被白白浪费掉。所以核心矛盾很清楚推荐序列本身是“变长”的而GPU只擅长“定长”的计算。解决思路无非两条——要么把变长序列改造成定长要么把计算方式改造成能容忍变长。业界最常见的方案就是“padding mask”也就是把序列补齐到统一长度同时用掩码告诉模型哪些位置是真实数据、哪些是补位。1.3 为什么不能直接放弃矩阵运算逐条跑有的同学可能会问既然GPU处理不规则序列这么麻烦那我干脆在CPU上逐条算或者GPU上开一个for循环遍历样本不行吗答案是能跑但效率极其惨烈。GPU适合并行计算你让它一条一条处理等于用跑车在胡同里倒车入库。具体来说逐条计算意味着每个batch只用一个GPU核心在工作其余几十上百个核心全部闲置。实际测下来一个batch 64条样本如果逐条算GPU利用率可能连5%都不到比CPU还慢——因为CPU的单核性能通常比GPU单核强而且省掉了数据从CPU拷贝到GPU的传输开销。还有一个更隐蔽的问题推荐系统里的模型往往非常重。以主流的深度兴趣网络DIN、兴趣演化网络DIEN为例除了序列编码层还有大量的Embedding层、全连接层、注意力计算。这些计算本身对数据形状非常敏感形状不一致就意味着没法走高度优化的cuBLAS或者cuDNN算子只能回退到框架的通用算子实现性能差距可能是一到两个数量级。所以把不规则序列“改造”成GPU擅长的规则计算不是可选项而是把模型训练和推理跑快的必答题。2. 改造方案设计Padding、Mask 与 Batch 构建的组合拳2.1 Padding 不是简单补零位置信息得提前想清楚把不规则序列变成规则矩阵最直接的做法就是Padding——选定一个最大长度max_len所有样本的序列都截断到这个长度不足的部分用特殊值填充。但这里有个新手很容易踩的坑padding值的选择会直接影响模型的学习效果。如果序列特征是Embedding的索引比如商品ID、类目ID那padding位置一般用0填充并且在Embedding矩阵的第0行初始化成全零向量这样padding位置的Embedding就是零向量不会对后续计算产生有效信息。这个方案工程上最省事也是大多数开源框架的默认做法。如果序列特征是连续值比如价格、评分、时间戳就不能简单补零了。原因在于补零会引入“这个位置的数值就是0”的虚假信息。举个简单的例子用户对商品的评分如果是0~5的浮点数padding位置补0模型可能会学到“0是一种很低的评分”但实际这个位置根本不是有效评分而是空白。这种信息污染在训练数据充分的时候影响不大但在样本量有限或者特征稀疏的场景会造成可观测的性能回退。正确的做法是特征本身归一化到一个不含0的区间比如1~2然后padding位置补-1这样模型在归一化层就能明确区分有效位置和padding位置。或者干脆在生成batch时记录一个序列长度向量靠mask来屏蔽无效位置特征值本身随意但必须保证mask用对。截断策略同样值得重视。线上真实场景中用户行为序列的分布往往是长尾的大部分人只有几次行为但头部用户可能有几千次行为。如果max_len设得过小头部用户的历史信息大量丢失如果设得过大大部分人又被大量padding浪费显存和算力。我通常的做法是统计训练集序列长度的95分位或者99分位把这个值作为max_len的参考基准。比如统计出来99%的样本长度都在128以内那就把max_len设成128只牺牲1%的头部信息换取绝大部分样本的紧凑表示。2.2 Mask 机制告诉模型哪些位置说的是人话Padding只是把序列变成规则的形状但模型并不知道哪些位置是真实数据哪些是补出来的。如果直接拿padding后的序列去算注意力或者池化padding位置就会作为“噪音”参与计算污染最终的表征。Mask机制就是为了解决这个问题。它的本质是一张和输入同等形状的布尔矩阵或者0/1矩阵有效位置为1padding位置为0。Mask在整个网络里的传播方式可以类比为“人体围栏”——真实信息可以在里面自由流动补位信息被物理隔离。具体到不同类型的网络层Mask的用法不太一样Embedding层padding索引对应的向量初始化为全零这其实是一种最朴素的mask。算加法或者均值池化的时候零向量加进去不改变结果相当于自动被忽略了。但如果做的是sum池化而不是mean池化零向量的存在会让结果偏小需要在池化时根据实际长度做除法。Attention层注意力分数要加上一个极大的负偏置比如-1e9这样softmax之后padding位置的权重就趋近于0。这一步必须做否则softmax会把padding位置的key也算进注意力加权和里。池化层无论是max池化还是mean池化都需要先用mask把padding位置替换成极小值max池化或者对非padding位置求和再除以真实长度mean池化。Mask还有一个容易忽略的作用在计算loss的时候如果样本的标签本身就对应序列的某个位置比如下一个商品预测那padding位置对应的预测结果必须从loss里mask掉否则模型会努力去“预测空白”白白浪费模型容量还可能训练出不稳定的梯度。2.3 Batch 构建的工程细节从 collate_fn 到内存布局数据从原始样本到GPU上的规则张量中间要经过数据加载器DataLoader的collate_fn。这个函数是改造链路里最容易出性能瓶颈的地方也是代码上相对容易优化的一环。最常见的写法是在collate_fn里对每条样本做padding然后把padding后的结果转成torch.Tensor最后堆叠成一个batch。这种方式逻辑清晰但有一个问题如果数据处理里有Tokenizer或者连续值归一化这类耗时操作逐条处理的开销会很大。更高效的做法是在Dataset的__getitem__里只做轻量级的预处理把原始序列连同长度信息一起返回真正的padding放到collate_fn里统一处理。另外一个细节是padding之后的内存布局要尽可能连续。PyTorch里torch.stack比torch.tensor(list_of_arrays)通常更快因为stack可以直接复用底层分配好的连续内存而tensor构造函数需要做一次类型推断和复制。实测下来在样本长度维度不大的情况下两种方式差异不大但当序列很长、特征维度很高时stack的优势会明显体现。如果追求极致性能还有一个高级操作把同一个batch内的样本按照长度降序排列先处理长样本。这主要是为了配合某些框架的循环神经网络RNN实现让模型可以提前终止对padding位置的计算。不过对于纯Transformer架构或者池化方案排序带来的收益有限可以不考虑。数据加载还有一个“隐形坑”DataLoader默认的num_workers是0也就是在主进程里加载数据很容易拖慢GPU的利用率。我自己的经验是在CPU核数允许的情况下把num_workers设成8或者16同时把pin_memoryTrue打开让数据在CPU端就固定到页锁定内存减少H2D拷贝的耗时。这两项改动几乎是零成本但能让训练吞吐提升10%~30%。2.4 工具选型解析直接用框架内置能力还是自己造轮子在改造不规则序列这件事上业界已经有了一些成熟的工具不一定要完全从零手写。如果你用的是PyTorchtorch.nn.utils.rnn.pad_sequence和pack_padded_sequence/pad_packed_sequence是处理变长序列的官方方案。其中pad_sequence用来把一列长度不等的Tensor补成规则矩阵pack_padded_sequence可以把padding后的序列“压缩”回变长形式让RNN在计算时跳过padding位置。这两个函数组合起来是PyTorch里处理变长序列的标准姿势。如果你用的是TensorFlow/Kerastf.keras.preprocessing.sequence.pad_sequences提供了类似的能力配合Masking层可以自动忽略padding位置的更新。但这里有一个重要的判断标准你是需要“在训练时动态处理变长序列”还是“在推理时快速处理固定长度的输入”。如果是前者框架自带工具基本够用如果是后者尤其是线上推理服务对延迟很敏感建议把padding逻辑固化到数据处理管线里提前把序列补到固定长度推理时直接跑batch预测省掉动态padding的开销。还有一类场景需要自己写代码序列里每个位置不止一个特征而是多个特征拼接比如商品ID 类目ID 价格 时间戳。这时候不能直接调pad_sequence因为每个位置的特征维度不一样。我的做法是先对每个特征列分别做padding得到一组形状相同的矩阵然后在最后一维上做concat。注意padding索引对齐即可concat本身不影响mask的正确性。3. 实操过程与核心环节实现从数据处理到模型代码的逐步落地3.1 环境准备与基础依赖我开始改造之前先确认手里的环境是否完备。因为后续要大量使用GPU算力驱动和CUDA版本必须和深度学习框架匹配。我的环境是Ubuntu 20.04、NVIDIA驱动470系列、CUDA 11.4、PyTorch 1.12.0。不同版本的PyTorch对CUDA版本有明确要求建议先跑一句python -c import torch; print(torch.cuda.is_available())确认GPU可用。如果你用的是更新的PyTorch版本比如2.xCUDA版本要求可能更高安装的时候注意看清楚安装命令里的cu118、cu121这类后缀。我个人踩过最大的坑是驱动版本够新但安装的PyTorch编译时面向的CUDA版本和驱动不兼容导致torch.cuda.is_available()永远是False。数据方面我这边用了一个模拟的用户行为数据集。每条样本的结构是user_id整数ID表示用户item_seq整数列表表示用户按时间排序的点击商品ID长度从1到100不等item_seq_len整数表示item_seq的实际长度label0/1表示这次会话是否有转化数据量不大一万条主要是方便调试。3.2 数据预处理的完整代码实现先定义一个简单的Dataset类里面负责读取样本、做一些必要的格式转换。import torch from torch.utils.data import Dataset, DataLoader import numpy as np import random class UserBehaviorDataset(Dataset): def __init__(self, samples, max_len128): self.samples samples # 预计算一个阈值超过max_len的截断不超过的保留原长度 self.max_len max_len def __len__(self): return len(self.samples) def __getitem__(self, idx): sample self.samples[idx] item_seq sample[item_seq] seq_len len(item_seq) label sample[label] # 这里只做轻量级预处理 # 截断逻辑放在collate_fn里统一处理会更好但为了演示清晰先在这里截断 if seq_len self.max_len: item_seq item_seq[-self.max_len:] # 保留最近的行为丢弃最早的 seq_len self.max_len # 返回Python原生的list和整数padding放到collate_fn return { item_seq: item_seq, seq_len: seq_len, label: label }注意我在__getitem__里返回的是list而不是已经padding好的Tensor。这是为了让padding操作集中在collate_fn里避免每条样本都做一次无谓的Tensor创建。接下来是collate_fn这是改造的核心战场def collate_fn(batch): # batch是一个list每个元素是Dataset.__getitem__的返回值 item_sequences [item[item_seq] for item in batch] seq_lengths [item[seq_len] for item in batch] labels [item[label] for item in batch] # 找到当前batch里的最大长度做动态padding # 用动态长度而不是全局max_len可以节省显存 batch_max_len max(seq_lengths) padded_seqs [] for seq in item_sequences: # 先在末尾补0补到batch_max_len pad_len batch_max_len - len(seq) padded_seq seq [0] * pad_len # 0作为padding索引对应item_id0我们在embedding层让0的向量全零 padded_seqs.append(padded_seq) # 转成Tensor item_seq_tensor torch.tensor(padded_seqs, dtypetorch.long) # 形状: [batch_size, batch_max_len] # mask矩阵有效位置为1padding位置为0 # 直接用seq_len生成避免逐元素比较 mask torch.zeros_like(item_seq_tensor, dtypetorch.bool) for i, seq_len in enumerate(seq_lengths): mask[i, :seq_len] True label_tensor torch.tensor(labels, dtypetorch.float32) seq_len_tensor torch.tensor(seq_lengths, dtypetorch.long) return { item_seq: item_seq_tensor, seq_len: seq_len_tensor, mask: mask, label: label_tensor }这里有几个关键点第一padding索引选了0这意味着Embedding矩阵的第0行必须是全零向量。如果直接用PyTorch的nn.Embedding默认会随机初始化你需要手动把第0行的权重置零或者让模型保存一个额外的padding_idx0参数PyTorch会帮你把该行的梯度永远置零。第二mask矩阵不用遍历所有元素去比较而是根据seq_len直接填充。这样写不仅代码简洁而且效率高——生成一个全零矩阵再按行赋值局部True比逐元素判断快得多。第三这里每个batch采用动态max_len而不是全局max_len。好处是当batch里恰好都是短序列时计算的矩阵更小显存占用和计算量都更少。代价是每个batch的形状不同框架底层可能会有额外的shape切换开销。实际测试下来还是动态长度更划算。如果你希望代码更简洁可以直接用torch.nn.utils.rnn.pad_sequencefrom torch.nn.utils.rnn import pad_sequence padded_seqs pad_sequence( [torch.tensor(seq, dtypetorch.long) for seq in item_sequences], batch_firstTrue, padding_value0 )效果是一样的但你要额外创建Tensor再让pad_sequence内部处理。上面的手动实现其实更好理解也方便以后扩展成多特征拼接的场景。3.3 模型内部如何正确消费 Mask数据已经变成规则的Tensor了接下来是模型侧的改造。一个典型的序列编码模型结构是Embedding层 - 序列编码层可以是GRU也可以是Self-Attention甚至就是Mean Pooling - 全连接层 - 输出。import torch.nn as nn import torch.nn.functional as F class SequenceEncoder(nn.Module): def __init__(self, vocab_size, embed_dim64, hidden_dim128): super().__init__() # padding_idx0PyTorch自动把第0行的梯度置零不会更新它 self.embedding nn.Embedding(vocab_size, embed_dim, padding_idx0) # 序列编码层用GRU作为示例 self.gru nn.GRU( input_sizeembed_dim, hidden_sizehidden_dim, batch_firstTrue, bidirectionalFalse ) def forward(self, item_seq, mask): # item_seq: [batch_size, seq_len] # mask: [batch_size, seq_len] # 1. Embedding层 embed self.embedding(item_seq) # [batch_size, seq_len, embed_dim] # 2. GRU编码 # 这里直接用pad_packed_sequence来让GRU跳过padding位置 lengths mask.sum(dim1).cpu() # 因为pack_padded_sequence要求lengths在CPU上且降序 # GRU要求序列按长度降序排列所以我们先把序列按长度排序 sorted_lengths, sort_idx lengths.sort(descendingTrue) embed_sorted embed[sort_idx] packed_input nn.utils.rnn.pack_padded_sequence( embed_sorted, sorted_lengths, batch_firstTrue, enforce_sortedTrue ) packed_output, hidden self.gru(packed_input) # hidden: [1, batch_size, hidden_dim] # 取最后一个时刻的隐状态 output, _ nn.utils.rnn.pad_packed_sequence( packed_output, batch_firstTrue, total_lengthitem_seq.size(1) ) # 把排序后的顺序还原回来 _, unsort_idx sort_idx.sort() hidden hidden.squeeze(0)[unsort_idx] return hidden这段代码里有几个必须注意的细节第一pack_padded_sequence要求lengths降序排列。这是PyTorch的硬性要求。如果你不排序要么报错要么得到完全错误的结果。我习惯用sort做一次索引变换再用sort_idx.sort()反推原始顺序确保输出和输入对齐。第二lengths必须放在CPU上。如果mask.sum(dim1)结果是GPU上的Tensor直接传给pack_padded_sequence会报错。我上面先调用了.cpu()这一步很容易漏掉。第三pad_packed_sequence里最好指定total_length。因为不同batch的序列长度可能不同如果你在同一个模型里跑了多个batch输出形状可能不一致。指定total_lengthitem_seq.size(1)可以保证输出形状固定为[batch_size, max_len, hidden]方便后续操作。如果你不是用GRU而是用Self-Attention做序列编码那更简洁直接乘上mask矩阵的扩展版本就行class SelfAttentionEncoder(nn.Module): def __init__(self, embed_dim64, num_heads4, hidden_dim128): super().__init__() self.attn nn.MultiheadAttention( embed_dimembed_dim, num_headsnum_heads, batch_firstTrue ) self.linear nn.Linear(embed_dim, hidden_dim) def forward(self, item_seq, mask): # item_seq: [batch_size, seq_len, embed_dim] # mask: [batch_size, seq_len]True表示有效 # 注意力要求传入key_padding_maskTrue表示“需要被忽略”的位置所以取反 key_padding_mask ~mask # [batch_size, seq_len] attn_out, _ self.attn(item_seq, item_seq, item_seq, key_padding_maskkey_padding_mask) # 这里的attn_out形状和输入一致padding位置可能仍然有值但经过mask后这些值不会影响有效位置 # 再做一次mean pooling用mask把padding位置去掉 mask_expanded mask.unsqueeze(-1).float() # [batch_size, seq_len, 1] masked_attn attn_out * mask_expanded sum_vec masked_attn.sum(dim1) # [batch_size, embed_dim] len_vec mask.sum(dim1, keepdimTrue).float().clamp(min1.0) # 防止除零 mean_vec sum_vec / len_vec return self.linear(mean_vec)这里最关键的坑是nn.MultiheadAttention的key_padding_mask语义是“True表示需要被忽略”正好和我们定义的“True表示有效”相反所以必须取反。还有mean pooling时一定要除以真实长度mask.sum()如果直接用mean(dim1)padding位置的全零向量会把平均向量拉低。3.4 Loss 计算也要跟上 Mask模型输出之后loss的计算同样需要mask。尤其是序列预测这类任务假设你的模型输出的是一个长度也为seq_len的序列预测那每个位置的预测都有对应的标签。但padding位置根本没有标签必须把padding位置的loss置零。def compute_loss(logits, labels, mask): # logits: [batch_size, seq_len, num_classes] # labels: [batch_size, seq_len] # mask: [batch_size, seq_len] loss_fn nn.CrossEntropyLoss(reductionnone) # 返回每个位置的loss loss loss_fn(logits.permute(0, 2, 1), labels) # [batch_size, seq_len] masked_loss loss * mask.float() # 统计有效位置的loss均值 total_loss masked_loss.sum() / mask.float().sum().clamp(min1.0) return total_loss如果你用的是PyTorch自带的一些loss函数比如BCEWithLogitsLoss是不支持per-sample层面mask的。我一般用reductionnone这个选项自己再手动做mask可控性最强。3.5 DataLoader 与训练循环的完整连接最后把这些拼起来完整的训练循环如下from torch.utils.data import DataLoader def train_one_epoch(model, dataloader, optimizer): model.train() total_loss 0.0 for batch in dataloader: item_seq batch[item_seq].cuda() mask batch[mask].cuda() label batch[label].cuda() optimizer.zero_grad() output model(item_seq, mask) # [batch_size, hidden_dim] 或自定义输出 # 假设是二分类任务 logit model.fc(output).squeeze(-1) loss F.binary_cross_entropy_with_logits(logit, label) loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(dataloader) if __name__ __main__: samples load_samples() # 自行实现 dataset UserBehaviorDataset(samples, max_len64) dataloader DataLoader( dataset, batch_size32, shuffleTrue, collate_fncollate_fn, num_workers8, pin_memoryTrue ) vocab_size 10000 model SequenceEncoder(vocab_sizevocab_size) model model.cuda() optimizer torch.optim.Adam(model.parameters(), lr1e-3) for epoch in range(10): loss train_one_epoch(model, dataloader, optimizer) print(fEpoch {epoch 1}, Loss: {loss:.4f})这里有个小细节DataLoader的num_workers设成8pin_memory打开能让GPU的喂数据效率高很多。如果num_workers太大可能因为进程切换开销导致反而变慢建议根据CPU核数先设一个合适的值再微调。4. 常见问题与排查技巧实录4.1 GPU 利用率上不去为什么训练还是慢很多人遇到的一个现象是GPU利用率在nvidia-smi里看着有80%以上但训练速度就是上不去。排查下来最常见的原因是CPU数据加载成了瓶颈。在推荐场景里数据本身的预处理往往很重读取原始文件、解析特征、计算特征交叉、做归一化……这些如果都在__getitem__里同步做CPU要干大量杂活。而GPU算得再快也要干等CPU喂数据。这时候num_workers没有调大或者pin_memory没开就会出现“GPU不是主力CPU才是”的尴尬情况。我解决这个问题的方法是把数据预处理尽量提前到离线阶段线上加载的时候只需要读必要的字段。__getitem__里只做最轻量的转换比如list转Tensor把耗时操作全部放到数据准备阶段。适当调大num_workers给足CPU线程。如果CPU已经满载但GPU占用还是低那就得看是不是模型本身的计算量太小小而快的batch完全喂不饱GPU。这种情况可以考虑增大batch size或者用梯度累积来模拟更大的batch。4.2 显存 OOM但不是因为模型太大训练推荐模型时遇到OOM很多人第一反应是模型参数太大要去换小模型。但在处理不规则序列的场景OOM的元凶往往不是模型本身而是动态padding时某个batch里出现了一条特别长的序列。假设你的max_len设的64但某个batch里恰好有一条长度为200的样本比如头部用户动态padding时这个batch会被pad到200的长度别的样本都跟着遭殃显存瞬间翻了三倍。解决方案在Dataset里对超长样本做强制截断这是最可靠的。或者在collate_fn里设置上限如果某个batch的max_len超过阈值就直接截断长样本到阈值。还可以按照长度对样本做分桶bucket每个桶内部长度接近分桶后再组batch这样动态padding的长度不会因为个别样本而失控。我实际的项目里用了分桶方案长度0~10的样本在一个桶11~30在一个桶31~64在一个桶每个桶各自padding到桶内最大长度。这样batch内长度方差小不需要全局max_len显存和算力都能压到最低。4.3 模型训练得很好但线上推理结果不对这是一个特别容易踩的坑训练时用动态padding mask但线上推理时为了追求速度直接把所有样本pad到固定长度结果发现模型效果崩了。原因是训练时模型见过各种长度的输入包括短序列也学会了利用mask区分有效位置。但线上推理时如果把所有样本都pad到固定长度模型可能过度依赖位置信息比如位置编码导致短样本的表现和训练时大相径庭。正确的做法是推理时的数据处理流程要和训练时保持一致。要么也动态padding线上服务如果延迟允许完全可以动态pad要么训练的时候就统一用固定长度训练让模型自适应这种输入分布。最忌讳的就是训练一套逻辑、推理另一套逻辑。4.4 训练时“正常”但梯度爆炸在处理序列模型时梯度爆炸的问题比普通模型更常见尤其是在GRU/LSTM这类循环结构上。序列越长时间步反向传播的路径越深梯度越容易指数级放大或消失。如果你发现loss在某个epoch后变成NaN或者突然跳到极大值可以先检查一下学习率是不是太高然后确认序列编码层的梯度是否异常。我是这样排查的for name, param in model.named_parameters(): if param.grad is not None: print(name, param.grad.abs().max().item())如果某个参数的梯度max值达到几万甚至几百万那就是梯度爆炸了。缓解手段给优化器加clip_grad_norm_torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0)效果立竿见影。降低学习率。换用更稳定的序列编码器比如从GRU换成Transformer或简单池化。不过我还要提醒一点有时梯度爆炸是因为padding位置的mask没做对导致padding位置的embedding参与了反向传播给了模型一些奇怪的梯度信号。这种时候查mask就能发现“漏网”的问题。4.5 batch 排序导致模型输入不一致我们在GRU的实现里做了batch内部的长度排序这一步是为了满足pack_padded_sequence的要求。但排序有可能导致一个问题同一份数据因为batch内样本分布不同排序后的顺序不同影响模型的随机性。这本身不是bug但如果你把模型保存/恢复、做分布式训练要留意排序逻辑是否在每个rank上保持一致。一个更简单的方法是在Dataset级别就把样本按长度排序或者分桶。这样每个batch内部的长度天然接近pack_padded_sequence需要的排序变化不大甚至不用排序。我后来把分桶逻辑推到Dataset层之后collate_fn里的排序步骤直接省掉了大半代码也更清晰。5. 性能调优与扩展思路把改造后的计算推到极致5.1 Batch Size 和 Max Length 的协同调优做GPU加速batch size是最敏感的超参数之一。batch太小GPU利用率上不去batch太大显存容易OOM。但在不规则序列的场景你真的不能说“batch size 64”就完事了因为每个样本的实际长度可能差别很大。我通常用“总token数”而不是“样本数”来衡量一个batch有多大batch_total_tokens sum(seq_len_i)。以显存占用估算来说注意力层的显存复杂度是O(batch_size * seq_len^2)所以当batch里都是长序列时实际占用的显存可能比短序列大一个数量级。在实践中我倾向于对显存比较宽裕的场景比如A100 40GB用较大的batch但限制每个batch的总长度total_length_threshold。对显存较小的场景比如2080Ti 11GB宁可batch小一点也不要让长序列撑爆显存。当然这种“总长度阈值”需要你自己实验几次才能定下来。我一般先跑一个batch用torch.cuda.max_memory_allocated()记录峰值显存再调整阈值保证峰值显存不超过显卡显存的80%。5.2 混合精度训练的收益与坑现在主流GPU都支持FP16混合精度训练。在推荐序列场景混合精度可以显著加速计算减少显存占用。但注意不是所有层都能直接用FP16跑好的。我建议Embedding层保持FP32。因为embedding是查表操作精度敏感换成FP16容易导致embedding向量的信息损失。序列编码层GRU/MHA可以试试FP16但要注意梯度缩放。Loss计算和最终输出层保持FP32避免精度问题。PyTorch的torch.cuda.amp已经封装好了GradScaler和autocast用法很简单。我唯一踩过的坑是某些自定义的mask操作在FP16下会出数值问题比如把mask转换成float16类型之后用-1e9做注意力偏置半精度下-1e9可能被截断成-65504导致softmax分母溢出。解决办法是mask偏置一般用FP32或者在mask处显式指定dtypetorch.float32。5.3 从训练到推理TensorRT 与 ONNX 的适配如果你最终要上线那么在推理阶段用PyTorch直接跑动态padding会慢一些。可以考虑把模型导出成ONNX再用TensorRT加速。但这里有个坑ONNX导出时动态长度batch_max_len如果是一个动态维度TensorRT需要额外做动态shape profile配置否则导出的模型只能接受固定长度输入。因此很多工程实践里线上推理会退而求其次固定一个max_len把所有样本都pad到这个长度然后用TensorRT的静态shape优化。这又回到了前面说的“训练和推理一致性”问题。我的建议是如果线上打算用固定长度训练时也最好在固定的max_len下训练或者至少让模型见过固定长度的输入分布。这样推理时换到静态shape才不会出现效果回退。5.4 分布式训练下的序列特征处理如果你的数据集大到需要多卡或者多机训练序列特征的处理逻辑还需要额外关注两点第一每个GPU拿到的是一个mini-batch每个mini-batch里样本长度不同。如果你的collate_fn里用了动态padding那么每个rank上的batch形状可能不一样不影响训练收敛但会影响同步时的效率。常见做法是用DistributedSamplerpad_collate系列函数保证每个batch尽量均衡。第二Embedding表在分布式场景下可能太大需要做分片。推荐系统动辄亿级item的表单卡放不下。这时候用torch.distributed的ProcessGroup做embedding分片或者直接用DeepSpeed/ColossalAI这类框架的序列并行能力。但核心的mask逻辑不变只是数据的路由和聚合变复杂了。6. 一点私货改造过程中我沉淀的几条习惯整个项目做下来我最深的体会是不规则序列改造这件事写代码只是一小部分真正的难点在于把“数据形态”和“计算形态”对齐。最开始我做这个项目总觉得应该在模型里做所有事情所以在forward里写了很多复杂的mask逻辑结果代码可读性差、调试困难。后来我把padding逻辑尽量前移到数据加载阶段让模型根本看不见“不规则”这件事——输入进来的天然就是规则Tensor和配套mask模型只负责计算。这个思路转变之后代码简洁了很多bug也少了很多。还有一条经验是关于调试的。处理mask相关的问题最有效的工具不是打日志而是写“断言”。比如在collate_fn返回之后我会立刻检查assert mask.sum(dim1).min() 0, 存在空序列样本 assert (item_seq 0).sum() (~mask).sum(), padding位置和mask位置不一致这类断言在开发阶段帮助极大几乎每次跑数据都能提前抓出边界条件问题。最后想说的是别看mask和padding这些操作听上去平平无奇实际上它决定了你的模型能不能充分利用GPU算力、能不能稳定收敛、能不能上线上不翻车。推荐系统里90%的序列建模工程问题归根结底都出在“数据形态”和“计算形态”的匹配上。把这一层想清楚并扎扎实实做好比在模型结构上堆花活要重要得多。
返回列表