#2739 feat: dist_muon offloading in megatron
原始 PR · 作者 Zhichenzzz · 合并时间 2026-08-29 09:16
为 Muon 优化器新增 NVMe 磁盘状态卸载,规避主机 OOM
值得精读。核心设计决策有三个:其一,不重写 DistOpt 分桶/写回逻辑,而是子类化 `ChunkedOptimizerStateOffloader` 只覆盖分配器,用最小改动复用 Megatron 既有的 chuncked restore、`assert_master_weights_resident` 和 prefetch 钩子;其二,利用文件页可回收性而不是简单扩内存,从内核语义上解决 OOM;其三,用 tensor 标记 + `msync` 与 Megatron checkpoint 协议对齐,避免全量重分配和不可归因的写回。这些模式对任何“把显存/内存压力转移到磁盘”的需求都有借鉴价值。
参与讨论