三年前我写过一篇 Jemalloc Profile 背后的数学原理, 纯粹是觉得它的采样设计有意思, 写完 tag 老老实实打了 JustForFun. 当时的心态很朴素: 这东西挺巧妙, 但我这辈子大概用不上. 三年后, Qwen RL 的规模眼看着越来越大, oom 时不时地崩了一下, 它被我从故纸堆里翻出来, 领了一份正式工作.
RL 规模上来了, 就不再是个算法问题, 而是个系统工程. 显存大家都盯得很紧, 这次说的是另一条线: CPU 内存. 一条轨迹从产生到被训练消费, 中间要倒不知道多少道手, 经过一堆 Python 对象和 native buffer; 这些东西单个看都不起眼, 堆进一个长跑的分布式链路里, 账就得重新算了. 内存曲线能看见在涨, 但它只负责说”在涨”, 不说是谁分配的, 更不说谁还攥着不放. 最气人的是, 等 OOM 了再去看, 进程已经没了, 现场直接蒸发, 连句遗言都没有. 所以需求很明确: 一个能跟着任务一直跑的 memory profile. 周期性留快照, 记录每块内存是哪个栈分配的, 释放时把账扣掉; 这样就算进程最后被杀了, 至少还能翻它生前最后一份记录, 跟之前的快照比一比, 找出一路在涨的那些分配栈. 要求说穿了就一个字: 便宜. 贵的东西大家只在怀疑出问题时才舍得开, 而内存问题恰好都是跑上一段时间才露头的.
memray 很好, 就是有点贵
第一个试的是 memray. Python 栈, native 分配, 火焰图, 该有的信息它都有, 看上去就是为这事生的. 临时抓一段确实好用; 但想让它跟着训练任务常开, 开销就有点吃不消了 —— Python 调用密集一点的负载, 开着它直接慢出一倍多. 训练可受不了这个.
给 memray 加采样: 采样赢了, 开销没输
我的第一反应很直接: 全量太贵, 那就采样. jemalloc 那套按分配字节做概率采样的账, 三年前刚算过, 还热乎, 照抄就行. 于是给 memray 实现了同款采样: 在展开 native 栈, 进入记录流程之前, 先把不抽中的分配事件过滤掉. 原生分配的热路径确实轻快了, 快了一个数量级. 我很满意, 直到拿它跑 Python 调用密集的负载: 端到端耗时基本没动. 这就有点尴尬了. 记录明明砍掉了绝大部分, 省下的时间呢? 我把实验一层层往下拆: 先把内存事件全扔了, 只留 Python profile; 再把 profile callback 换成一个除了 return 什么都不干的 C 空函数. 拆到这一步, 开销居然还在.
原因在 memray 知道”当前 Python 调用链”的方式上: 它挂着 CPython 的 profile 事件, 每发生一次函数调用就 push 一次, 每返回一次就 pop 一次, 持续维护一份影子栈. Python 代码只要还在跑, 这份栈就得一直喂, 跟你采不采样一点关系都没有. 更过分的是, 解释器在叫 callback 之前, 自己还得先取出当前 frame, 摆弄一遍引用计数 —— 所以 callback 哪怕是个空函数, 这份钱也省不掉.
一句话: 采样采掉的是”记录内存”的钱, 采不掉”跟着 Python 跑”的钱. 那能不能反过来, 平时完全不跟, 真抽中一次分配的时候, 再去把栈取回来?
抽中时再取栈, 但 GIL 等不得
scalene 就是这么干的: 内存采样触发之后, 才同步去取 Python 调用栈. 方向确认: 取栈成本只发生在采样时刻, 不用陪着 Python 跑全程. 但顺着它的实现往下看, 取 frame 之前先构造了一个 GIL guard, 里面调的是 PyGILState_Ensure(). 当前线程没拿着 GIL 的话, 就得在这儿等. 在普通代码里等 GIL 天经地义, 但在 malloc 里等就是另一回事了. malloc 可能发生在任何 native 扩展的任何角落, 进来的时候线程手上可能已经攥着另一把锁. 只要出现下面这个顺序, 就是死锁:
线程 A: 持有 native 锁 M
-> malloc 被抽中
-> 等 GIL
线程 B: 持有 GIL
-> 进入 native 扩展
-> 等锁 M
两边都在等对方先松手. 要说清楚, 这是我读源码推演出来的风险, 不是线上真炸过. 但一个要塞进 RL 任务里常年跑的东西, 我不想在 malloc 里埋一颗雷, 然后指望所有业务代码替我记着锁顺序.
从 faulthandler 偷师
什么时候一个进程最需要打印自己的 Python 栈? 快死的时候. faulthandler 干的就是这活: 崩溃现场, 环境已经不允许它像正常 Python 代码那样体面地工作了. CPython 的注释写得很坦白: 线程可能早就把 GIL 放了, 得从线程局部存储里把 thread state 掏出来, 之后沿着解释器内部的 frame 链直接往上爬. 全程不碰 GIL. 照这个思路, 我实现了不需要 GIL 的 Python frame 采集, 接进 jemalloc 的采样路径. 最终流程朴素得有点不好意思: jemalloc 决定这次分配采不采; 抽中了, 就地抓一份 Python 栈; 这块内存释放的时候, 把对应样本的账扣掉; 后台周期性 prof.dump 落盘. 采样, 分配释放匹配, 聚合, 导出, 全是 jemalloc 原来的机制, 我只补了 Python 来源这一层. 绕了一大圈, 干活的还是三年前那个 jemalloc.
当然, 直接读解释器内部状态是有边界的: 这是 best-effort 实现, 得跟着解释器版本走; 看到的也只是经过 allocator 的分配, 不是进程的全部内存. 但回答”这内存是谁分配出来的”这个问题, 够用了.
zero overhead, 还倒找一点
数字就不贴了, 枯燥. 说感受: 在真实 RL 任务上开着它跑, 端到端 step 耗时趴在原地没动, 还往下出溜了一点. 不是 profiler 有什么加速魔法. 是 LD_PRELOAD 上手的时候, 顺手把 libc 自带的 allocator 换成了 jemalloc, 而 jemalloc 的 malloc/free 本来就更快. 采样, 抓栈, dump 当然都有成本, 但架不住换进来的 allocator 跑得快, 对冲完还倒找.
到这个份上, 它才配常开: 不花钱的观测才配常驻, 常驻的观测才等得到问题.
不要给算法老师添麻烦
之前做数据流治理的时候, 我就一直坚持一个原则: 不要给算法老师引入额外的心智负担. 算法侧继续按业务语义组织对象, 基础设施侧通过可观测性和周期性治理, 维护系统的效率. 模型和算法天天在迭代, 今天合理的对象组织, 明天加个字段, 多一条消费链路, 就重了. 要求每位算法老师做实验之前, 先修一门 allocator 和 profiler, 这不现实, 也不应该. 我更愿意把观测做好: 哪条路径变重了, 哪类分配一直降不下来, 先让这些变化被看见, 再结合真实的业务语义去治理.
这次的 memory profile, 就是给 RL 框架的可观测性补上的一块. 至于三年前那些采样公式, 当时我觉得它们 JustForFun; 现在它们跟着每个 RL 任务上班打卡. 知识这东西学了不亏, 它只是在故纸堆里排队, 等一个 OOM.
