深度学习工程师秘而不宣的代码重构三原则

　　去年十一月份，我蹲在三号机房改一个PyTorch的Loss封装——不是调参，是把原本混在train_step里头的gradient clipping、label smoothing、logits校准全扒出来重写。当时空调坏了，冷凝水顺着风管往下滴，一滴一滴砸在UPS机柜顶盖上，咚、咚、咚，跟倒计时似的。

　　王磊坐我斜后方，左手捏着半块凉掉的韭菜盒子，右手还在键盘上敲。他忽然把耳机一把扯下来甩桌上：“你这个detach().clone()放错地方了。”我抬头看了眼他屏幕——他连梯度都还没画出来，纯靠print中间张量shape猜的。我说：“没放错。”他回：“你敢不敢把forward里那个repeat_interleave删了试试？”我没试。但三天后我自己删了，因为下游模块传进来的是B×1×D，repeat_interleave之后变成B×K×D，而K根本不是batch_size的整数倍……那个K是动态采样数，它有时候是7，有时候是0，有时候是NaN。我原以为它是离散整数，结果它在某次负样本过采样里，从torch.where里吐出来一个nan索引——就那么一次，就那么一个训练step，模型权重突然全部飞掉，lr_scheduler直接跪了，学习率缩到1e-12。这个我敢打包票。

　　那会儿窗外下着雨。不是毛毛雨，是那种上海初冬的湿冷雨，贴着玻璃往下淌，把对面楼灯光糊成一片黄晕。张薇推门进来递了杯热豆浆，杯底还印着“光明”两个字，她指尖有点抖，估计刚从冷冻区取完数据——她那天下午在搞一个clip-adapter的embedding对齐，说要用蒸馏loss替代cross-entropy，结果蒸馏出来的teacher logits里头夹着三个inf，还是float32精度下的inf。没人知道inf怎么进去的。我们查了十七遍数据预处理，最后发现是h5py读取时候的chunk cache默认用了np.float64，转float32时溢出……这种事儿不值当写进paper附录，但能让你凌晨两点在jupyter里翻三千行日志，边看边啃冷馒头。

　　我说的三原则，第一条就是：只重构正在流血的代码。不是“看着不爽”，不是“命名不够函数式”，不是“违反SOLID”。是它已经在跑测试集时漏了一个class的grad_norm，或者checkpoint加载时莫名丢了一层BN的running_mean。别的都是伪需求。可上次组会李哲举手说“我觉得DataLoader子类应该抽象出Pipeline接口”，我差点把笔记本合上走人——他那个Pipeline连tensor.device都没做检查，纯靠注释提醒用户“请确保device一致”。后来他在dev分支commit了一千八百行，然后在PR里加了个note：“已测通batch_size=4”。我回复：“测了三轮，第几轮爆的CUDA OOM？”，他回：“第三轮……但我改了pin_memory=False。”我回：“你改了，但没提。”他没回。我关了网页，桌角外卖盒盖子歪着，辣子鸡丁浮在褐色酱汁上，没动过。

　　第二条：重构必须带可观测断点，且断点要插在反向传播路径上。不是print，不是logger.info，是hook，是register_full_backward_hook，是backward_hook里的ctx.saved_tensors[0].norm().item()实时打点。你得亲眼看见梯度在哪个node开始塌缩，哪个node开始爆炸，哪个node开始静默。王磊上个月重构一个多任务loss，他说自己做了“充分单元测试”，结果线上infer的时候，某路分支输出全为0——debug三天才发现，他在compute_loss()末尾加了个with torch.no_grad():，忘了删。不是疏忽，是那个no_grad被他当成防止泄漏的保险丝了。他接了个电话，声音很小：“喂？哦…我这儿正卡在一个backward hook没触发……嗯，我再看看是不是hook挂错了地方。”挂了电话，他把咖啡杯往右推了十厘米，杯子底沿在键盘托盘上刮出一道白痕。

　　第三条……第三条其实挺难说清楚。它不是技术规范，是节奏感。比如你改完一个model.forward，别急着跑full val，先拿一个mini-batch，开profiler，看aten::add_耗时是不是比原来高了三倍；再拿同一batch走一遍eval mode，看dropout有没有偷偷生效；最后才跑完整epoch。张薇那天中午来问：“你重构的那个attention mask逻辑，为什么training loss震荡更小了？”我没回答。因为我重构的只是mask生成部分，真正让loss稳下来的，是顺手把mask tensor从cpu挪到了device上——但那是我第二天早上补的，和重构本身无关。所以我说：优胜劣汰是真的，但汰掉的往往不是烂代码，是烂时机。上周五下午三点四十七分，我删掉自己写的第三个版本refactor分支，因为scheduler.step()的位置移动后，在DDP sync之前触发了state_dict update，导致不同rank保存的optimizer state错位……这事儿我没法写进checklist，只能记在脑子里——哪天又下雨，机柜又滴水，我就想起这事。

　　对了，王磊现在用rust-tch重写他的loss module，说“编译期能抓住device mismatch”。我听着，没接话。他桌上的韭菜盒子早就干了，皮上裂出细纹，像某种神经元突触的简化图。

　　我昨天看到一段代码，里面有个变量叫output_final_just_for_debug_v2，后面跟着三行注释：“此变量不参与任何计算”、“此变量仅供vscode debugger inspection”、“不要删除，否则CI报错”。我不懂它为什么能过CI。也不打算去查。等它哪天开始影响梯度流，我再碰。

　　我现在在重读去年十一月的commit log，翻到那一行：“refactor loss compute —— remove repeated softmax, move clip to backward”。下面写着author: me，date: 2023-11-08 23:16:42。时间戳很准。但我记得那天其实是凌晨一点过五分，因为电梯里监控坏了一小时，我走消防楼梯下去，手机没电自动关机了。所以git用的是系统时间，不是我的真实时间。这个我说不准。

文章配图，仅供参考

（编辑：站长网）

【声明】本站内容均来自网络，其相关言论仅代表作者个人观点，不代表本站立场。若无意侵犯到您的权利，请及时与联系站长删除相关内容!