分子模拟异构算力适配开发教程（17）：推理调度思想迁移——从 PagedAttention 与 continuous batching 到 MD 作业调度器

分子模拟异构算力适配开发教程17推理调度思想迁移——从 PagedAttention 与 continuous batching 到 MD 作业调度器版本声明块工具/软件调度思想源自 vLLMarXiv:2309.06180SOSP 2023与 OrcaOSDI’22参考工程 Colmenaexalearn语言/环境Python 3.10调度器实现可对接第 14-16 篇的任意调度底座本文目标读完你能说清“为什么 MD 像推理服务”并跑起来一个带槽位回收与动态补位的 MD 批调度器一句话结论LLM 推理的continuous batching迭代级调度每次前向迭代重排 batch、完成的请求立即退出出自 OSDI’22 OrcaAnyscale 博客实测 23 倍吞吐思想可直接迁移到 MD——MD 作业同样是“短迭代循环 无相互依赖 吞吐导向”的负载把“每批等待最慢作业”的静态批static batching换成“槽位完成即补位”的迭代级调度单卡吞吐可数倍提升。〇、本篇要解决的认知问题PagedAttention 与 continuous batching 各解决什么问题两者的出处与证据是什么为什么说 MD 是天然的批调度负载它与 LLM 推理的同构点在哪MD 版的“批重排”调度器怎么设计槽位、回收、补位现有科学计算工作流系统Colmena/Swift/T与自研轻量调度器的边界在哪一、机制解析1.1 两个推理调度发明出处与机制为什么这一节对你重要你的团队大概率没做过推理服务但“推理调度封装”正是本主题异构算力适配任务书里的关键词——理解这两个机制的问题域它们为什么被发明才能判断迁移到 MD 是否成立。PagedAttentionvLLM 核心技术论文 arXiv:2309.06180SOSP 2023受操作系统虚拟内存与分页启发——KV cache 划分为 blockblock 如页、token 如字节、sequence 如进程逻辑 block 经 block table 映射到非连续物理 block按需分配、支持 copy-on-write 共享内存浪费 4%。它解决的是显存碎片与预留浪费传统方案按最大序列长度预留 KV cache。continuous batching迭代级调度机制出处是 OSDI’22 论文Orcaiteration-level scheduling——每次前向迭代重排 batch完成的请求立即退出、新请求立即加入“continuous batching” 这个名字由 Anyscale 博客《How continuous batching enables 23x throughput in LLM inference while reducing p50 latency》普及标题就是实测数字23 倍吞吐 p50 延迟下降。它解决的是静态批的队尾浪费——静态批要等整批里最长的序列生成完才能整批返回短请求陪着长请求空等。1.2 同构性论证MD 为什么像推理把两类负载的特征摆在一起维度LLM 推理MD 作业基本单元一次前向迭代一步积分帧作业形态迭代循环直至序列完成迭代循环直至 nsteps 完成作业间依赖请求相互独立各 MD 副本相互独立增强采样的定期交换除外——那是第 15 篇 gang 的事时长异质性序列长短差异巨大作业步数/体系大小差异巨大资源形态显存KV cache 算力显存坐标/力数组 算力优化目标吞吐token/s与延迟吞吐ns/day 总量与周转同构点的本质“迭代循环 作业独立 时长异质”三条全中。静态批的队尾浪费在 MD 上的等价物把 10 个 MD 作业打包到一张卡上排队前 3 个 20 分钟跑完、最后一个要 6 小时——静态调度下要么整批等槽位空转 5.5 小时要么整批撤重启浪费。continuous batching 的答案是槽位完成即回收、新作业即补位——卡上永远满负荷。异质点也要诚实列出迁移不是照抄MD 单作业启动成本高grompp/内存分配/首次 JIT没有推理那种“毫秒级增量”的优雅显存占用基本恒定体系定了就定了PagedAttention 解决的“KV cache 按需增长”问题在 MD 上不突出——所以值得迁移的是 continuous batching 的调度思想不是 PagedAttention 的内存机制后者对 MD 的对应物是 vGPU 切分第 14 篇但受铁律 9 约束。1.3 MD 版批调度器的设计要素从 1.1/1.2 推出的设计本篇代码的实现蓝图槽位slot模型一张卡 若干槽位整卡模式下 1 槽切分模式下 N 槽——铁律 9 实测说了算迭代级回收不盯“作业完成”这一终点而是周期性巡检作业完成/失败即回收槽位动态补位回收的槽位从优先级队列取下一个作业填上——不等“整批”概念优先级与公平交互作业用户在等优先于夜间批量——带 aging等待越久优先级越高防饿死幂等与断点补位重启的作业要有 checkpoint/续跑语义gmx 的 -cpi第 20 篇平台化。1.4 与现有工作流系统的边界不自研一切的判断依据真实项目调研底账Colmenaexalearn/colmena——“steer large campaigns of simulations on supercomputers”Thinker 决策 agent Doer 执行 task server 基于 Parsl论文 IJHPCA 2024适合“大规模模拟战役 中间结果驱动决策”的科研场景Swift/Tswift-lang/swift-t隐式并行 workflow 语言万亿次任务级适合超大规模任务图。边界两者都是任务编排框架不是单卡槽位级吞吐调度器——你要的“一张卡上多作业动态补位”这类细粒度逻辑要么下沉到底座K8s/HAMi 的 vGPU第 14 篇要么自己写一个几百行的轻量调度器本篇。第 20 篇的平台架构里本篇调度器是“应用层吞吐优化”Colmena 这类是“科研工作流层”——各居其位。二、完整代码与逐行剖析一个完整可跑的 MD 批调度器槽位回收 动态补位 优先级 aging 吞吐统计——本系列代码量最大的一篇也是第 20 篇平台的调度内核雏形#!/usr/bin/env python3MDScheduler推理调度思想continuous batching的 MD 版实现。 核心机制对照 iteration-level scheduling → 槽位完成即回收不等整批 dynamic batching → 回收即补位优先级队列 aging 防饿死 运行模式 demo用模拟作业无 gmx 依赖验证调度逻辑 gmx 真实调 gmx mdrun-nsteps 当作业时长槽位并发的 mdrun 进程 from__future__importannotationsimportjsonimportosimportrandomimportsubprocessimporttimefromdataclassesimportdataclass,fieldfromenumimportEnumclassState(Enum):QUEUEDqueued;RUNNINGrunning;DONEdone;FAILEDfaileddataclassclassMDJob:job_id:strtpr:str# gmx 模式tpr 路径demo 模式仅作标识nsteps:int# 作业时长demo虚拟秒数gmx-nstepspriority:int0# 基础优先级越大越优先交互作业给高值submit_ts:floatfield(default_factorytime.time)state:StateState.QUEUED start_ts:float|NoneNoneend_ts:float|NoneNoneproc:subprocess.Popen|NoneNone# gmx 模式进程句柄deadline:float|NoneNone# demo 模式虚拟完成时刻classMDScheduler:def__init__(self,slots:int2,mode:strdemo,aging:float0.01,poll:float0.2):self.slotsslots# 槽位数整卡1切分/多进程按实测定铁律 9self.modemode self.agingaging# 每等待 1 秒优先级加成防饿死self.pollpoll# 巡检周期iteration-level 的迭代粒度self.jobs:list[MDJob][]# ── 优先级 基础优先级 等待时间 × agingaging 机制──────────def_eff_priority(self,j:MDJob)-float:waitedtime.time()-j.submit_tsreturnj.prioritywaited*self.agingdef_pick_next(self)-MDJob|None:ready[jforjinself.jobsifj.stateState.QUEUED]returnmax(ready,keyself._eff_priority)ifreadyelseNone# ── 作业启动两种模式共用接口 ──────────────────────────────────def_start(self,j:MDJob)-None:j.state,j.start_tsState.RUNNING,time.time()ifself.modedemo:# demo虚拟时长nsteps 字段当秒数到点判完成j.deadlinetime.time()j.nstepselse:# gmx单卡多进程并发可见性隔离下都写 -gpu_id 0——第 14/16 篇的约定cmd[gmx,mdrun,-s,j.tpr,-deffnm,frun-{j.job_id},-nsteps,str(j.nsteps),-nb,gpu,-pme,gpu,-update,gpu,-gpu_id,0,-noconfout,-pin,on]j.procsubprocess.Popen(cmd,stdoutsubprocess.DEVNULL,stderrsubprocess.DEVNULL)# ── 巡检iteration-level 的核心——完成即回收、回收即补位 ─────────def_check(self,j:MDJob)-None:ifself.modedemo:iftime.time()j.deadline:j.state,j.end_tsState.DONE,time.time()else:rcj.proc.poll()ifrcisnotNone:# 进程退出0完成j.stateState.DONEifrc0elseState.FAILED j.end_tstime.time()defrun(self)-dict:runninglambda:[jforjinself.jobsifj.stateState.RUNNING]whileany(j.statein(State.QUEUED,State.RUNNING)forjinself.jobs):forjinrunning():self._check(j)# ① 巡检在跑的whilelen(running())self.slots:# ② 空槽即补位不等整批nxtself._pick_next()ifnxtisNone:breakself._start(nxt)time.sleep(self.poll)returnself.stats()defstats(self)-dict:done[jforjinself.jobsifj.stateState.DONE]span(max(j.end_tsforjinself.jobs)-min(j.submit_tsforjinself.jobs))ifdoneelse0return{total:len(self.jobs),done:len(done),failed:sum(j.stateState.FAILEDforjinself.jobs),wall_s:round(span,1),throughput:round(len(done)/span,3)ifspanelse0}defdemo()-None:演示8 个时长异质的作业 × 2 槽位对比静态批与连续批。random.seed(42)jobs[MDJob(fj{i},tpr-,nstepsrandom.uniform(2,12),priority5ifi2else0)foriinrange(8)]schedMDScheduler(slots2,modedemo)sched.jobsjobsprint(── continuous batching完成即补位──)resultsched.run()print(json.dumps(result,ensure_asciiFalse))# 静态批对照分 2 批每批 4 作业批内等最慢的# 简化模拟批耗时 max(作业时长) × 批数 / 槽位并发……durationssorted(j.nstepsforjinjobs)static_spansum(max(durations[i:i4])foriin(0,4))# 每批 4 个排队 2 槽print(f── 静态批估算 ── wall≈{static_span:.1f}s队尾浪费的代价)if__name____main__:demo()# gmx 模式MDScheduler(slotsN, modegmx) 真实 MDJob 列表逐段剖析_check 补位循环就是 continuous batching 的 MD 化Orca 的“每次迭代重排”在这里变成“每个 poll 周期巡检”——粒度不同推理是毫秒、MD 是秒级机制同构完成即回收、回收即补位。注释里的①②标出了这两个动作在循环里的位置。aging 机制_eff_priority优先级 基础 等待 × 系数——交互作业priority5先走但批量作业等的越久分越高最终必被调度防饿死的经典解法。aging0.01 意味着等 500 秒的批量作业等效一个交互作业——系数按业务 SLA 调。demo 模式的意义调度逻辑的验证不需要 gmx/GPU——CI 里秒级跑通“完成即补位”的时序gmx 模式只是_start/_check的另一个实现进程句柄替换虚拟时钟——调度内核与执行后端分离第 18 篇的适配层会把这个模式推到引擎级。静态批对照的static_span计算两批各自等最慢成员——8 个作业里 12 秒的那个让同批三个短作业陪着空等这就是 1.1 节“队尾浪费”的可计算形态。跑 demo 会看到 continuous batching 的 wall 显著小于静态批估算差距随时长异质性增大。三、常见报错与排查问题 1现象——gmx 模式下单槽多进程并发跑ns/day 全体暴跌。根因槽位并发进程数超过卡的承载——N 个 mdrun 争一张卡的算力与显存带宽MD 带宽敏感铁律 9 的进程级版本另外显存超限会 OOM 部分进程。解法槽位数用第 12 篇流水线实测标定1 进程 vs 2 进程的各自 ns/day × 数量 总吞吐取最大点体系大的作业标记独占slots 约束为 1或走 HAMi vGPU 显存隔离第 14 篇但要过铁律 9 实测。问题 2现象——补位重启的作业从头跑前面几小时的采样白费。根因调度器只管进程生死不管作业的续跑语义——mdrun 被杀/退出后没带 checkpoint 续跑。解法作业模板统一带-cpo/-cpiGROMACS 的 checkpoint 续跑具体选项语义以 gmx mdrun 在线帮助为准与调度器的重启钩子衔接第 20 篇平台把“断点续跑”做成适配层的标准能力幂等性的一部分。问题 3现象——调度器自己成了瓶颈巡检循环吃满一个核。根因poll 周期太小毫秒级轮询或作业数巨大时_pick_next每轮全表扫描。解法poll 按业务粒度设MD 作业秒级完成差异无意义0.2–1 秒足够作业量大时把 ready 队列换成堆heapq按有效优先级——本篇的线性扫描在百级作业内没有问题千级再优化。问题 4现象——优先级反转高优交互作业提交后仍排在某个长作业后面。根因交互作业的 priority 忘了给高值或 aging 让老批量作业的分超过了它也可能是槽位全被长作业占着补位要等回收。解法交互作业 priority 显式给值demo 里 i2 给 5 的模式槽位全占时可设“抢占预留”最后一个槽位只给高优作业——这是平台 SLA 策略问题第 20 篇的配置面会暴露这个开关。四、动手练习练习 1基础跑 demo()记录 continuous batching 与静态批估算的 wall 对比再改random.uniform(2, 12)为random.uniform(5, 6)时长同质化观察差距消失。判定成功标准两组数字都拿到能解释为什么时长异质性越大、连续批优势越大对照 Orca 的 iteration-level 语义同质化后两种策略 wall 接近差距 10%。练习 2进阶给调度器加统计每作业的排队时长start_ts - submit_ts与“年龄最大时的优先级”输出 top-3 等待作业。判定成功标准统计输出正确排队时长 ≥0 且和 wall 时间自洽验证 aging 生效批量作业的有效优先级随时间上升交互作业排队时长显著小于批量作业。练习 3思考题无标准答案如果把 PagedAttention 的“分页”思想硬搬到 MD会对应什么可行吗思考方向验证要点① MD 的“页”应该是什么帧残基块的坐标——坐标数组的分页对显存利用有什么影响② KV cache 的 copy-on-write 共享在 MD 的对应物增强采样里共享的初始构象③ 为什么说“MD 显存基本恒定”让这个迁移价值有限——什么样的负载如在线采样的体系增长才会让它复活。五、小结与下一篇预告本篇完成了推理调度思想的迁移论证与实现PagedAttention显存分页4% 浪费与 continuous batching迭代级批重排Orca 出机制、Anyscale 出 23 倍实测是两个不同层面的发明MD 与推理的三条同构迭代循环/作业独立/时长异质支撑了 continuous batching 的迁移——队尾浪费的 MD 版就是“多作业排队等最慢成员”调度器的三要素槽位回收/动态补位/优先级 aging在 MDScheduler 里完整落地demo 模式把验证成本降到零依赖。Colmena/Swift/T 划定了“自研轻量调度器”的合理边界。下一篇做适配层统一 GROMACS 与 OpenMM 的 Python 封装——能力探测平台枚举/gmx 版本解析、统一作业描述、后端协商、运行时环境注入CUDA_VISIBLE_DEVICES 的三层可见性在这里收口。它是本篇调度器与第 20 篇平台之间的桥。本篇认知问题回显FAQQ1continuous batching 是什么出处是哪里A指推理服务的迭代级调度——每次前向迭代重排 batch完成的请求立即退出、新请求立即加入机制出自 OSDI’22 论文 Orcaiteration-level schedulingcontinuous batching 这个名称由 Anyscale 博客普及并实测 23 倍吞吐提升它解决静态批的队尾浪费整批等最长序列。Q2为什么说 MD 作业适合 continuous batching 式调度A因为 MD 与 LLM 推理在三条负载特征上同构基本单元是迭代循环积分步/前向步、作业间相互独立各 MD 副本无依赖、时长异质性大步数与体系规模差异大——静态批的队尾浪费在 MD 上表现为多作业排队等最慢成员槽位完成即回收补位可显著提升单卡吞吐。Q3MD 版批调度器需要哪些核心机制A四件槽位模型整卡 1 槽或按实测的切分多槽、迭代级回收周期巡检完成/失败即回收槽位、动态补位回收立即从优先级队列取新作业、优先级 aging有效优先级基础优先级等待时间×系数防饿死配套断点续跑gmx checkpoint保证补位重启不浪费已有采样。Q4自研 MD 调度器和用 Colmena/Swift-T 的边界在哪AColmenaexalearnThinker 决策Doer 执行基于 Parsl面向大规模模拟战役与中间结果驱动决策的科研工作流Swift/T 面向万亿次任务级的隐式并行 workflow——两者都是任务编排层单卡槽位级的吞吐优化完成即补位、显存切分协调要么下沉给 K8s/HAMi要么用几百行的轻量调度器自研本篇 MDScheduler第 20 篇平台中两者分层共存。