Repository navigation
fix(megatron): repack oversized KK microbatches - #419
DrRyanHuang wants to merge 1 commit into
Conversation
# 🐛 Bug Fix ## Fix dynamic microbatch packing - Return actual sample-index groups from First-Fit packing - Fall back to First-Fit when KK groups exceed the token budget - Split fallback groups to match DP/VPP microbatch counts
|
Thanks for contributing to Relax, @DrRyanHuang! 感谢你为 Relax 做出贡献！ Contribution guide / 贡献指南Describe the problem, your changes, and how you validated them. Keep each PR focused and run 请说明问题、改动和验证方式，保持 PR 聚焦，并在提交前运行 CI commands / CI 指令
Put one command on the first line of a new PR comment. Rerun/cancel require PR authorship or repository write access. 在新 PR 评论的首行写一条指令。PR 作者或有仓库写权限的贡献者可以重跑、取消 CI。 |
Nyanpasu 审查看板审查状态： 💬 已完成 · 有补充意见 审查版本： 11edf3c 目标分支： main 审查完成，结论为 Comment：打包回退机制本身经验证正确（计数等价 3000 组、回退正确性 6000 组随机用例），无阻塞问题；提出 1 个 P2 测试缺口（F1）与 1 个 P3 非行级跟进建议（F2，SFT 路径同样缺少 KK 容量上限）。
审查发现待处理
已解决或已取代
提交范围 · 接收 2 · 建议移出 0 · 待确认 0接收 2 个文件 · 建议移出 0 个文件 · 待确认 0 个文件。移出与待确认部分暂停深审，不代表审查通过。
精简审查与验证依据
生产代码的必要性与替代方案
测试的必要性与替代方案
Powered by Nyanpasu with glm-5.3-flash max, please check the suggestions carefully.
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
rai-studio-bot
left a comment
rai-studio-bot
left a comment
There was a problem hiding this comment.
审查结论
本 PR 的修复经验证正确：重构后的 get_minimum_num_micro_batch_size 与原实现计数逐例一致（3000 组随机用例）；对真实的 KK 分组与新的 First-Fit 回退（含按需拆包）做了 6000 组随机用例模拟，在 DP/VPP 包数对齐不变量下，回退后每个样本恰好覆盖一次、多样本分包均不超过 token 预算、包数与 dummy offset 均正确。现有调用方 relax/backends/megatron/actor.py:1461 的计数语义不受影响，pre-commit 全部通过。
看板（范围结论与验证依据）：#419 (comment)
-
relax/backends/megatron/actor.py:1462（以及:1294的 DP 对齐重分包）同样先以 First-Fit 计数、再用get_seqlen_balanced_partitions分组，且没有容量检查，与本 PR 修复前的 RL 路径模式相同。若 SFT 的 pack 长度可能严重不均，会有同样的超预算风险。本 PR 明确限定 RL 路径，此点留作后续跟进即可。
| if any(sum(seqlens[index] for index in partition) > capacity for partition in partitions): | ||
| # KK balances load without a capacity constraint. Only split the | ||
| # first-fit bins when DP/VPP require more micro-batches. | ||
| partitions = get_first_fit_partitions(seqlens, capacity) | ||
| while len(partitions) < real_partition_count: | ||
| partition = max(partitions, key=len) | ||
| partitions.append([partition.pop()]) |
There was a problem hiding this comment.
回退逻辑本身经验证是正确的：用真实的 get_first_fit_partitions 与 get_seqlen_balanced_partitions 加上本函数原文随机模拟 6000 组用例（含 DP/VPP 包数对齐不变量），回退后每个样本恰好覆盖一次、多样本分包均不超过 capacity、包数与 dummy offset 均正确；get_minimum_num_micro_batch_size 重构前后计数也逐例一致。
但目前没有测试保护这条路径：tests/backends/megatron/test_data_vpp.py 只覆盖 use_dynamic_batch_size=False，tests/utils/data/test_seqlen_balancing.py 只测 KK 分组。这段回退分支与 get_first_fit_partitions 都是纯函数，不依赖 GPU/分布式即可直接测试。建议补一个回归测试，至少覆盖：KK 超预算触发回退、回退后各分包不超过 capacity、样本覆盖完整、len(partitions) < real_partition_count 时按需拆包到目标包数。否则后续重构可能悄悄退回超预算打包，重新引入 OOM 风险。
| partitions = get_first_fit_partitions(seqlens, capacity) | ||
| while len(partitions) < real_partition_count: | ||
| partition = max(partitions, key=len) | ||
| partitions.append([partition.pop()]) |
There was a problem hiding this comment.
看起来就是 KK 超预算时候完全回退到 FF？话说是否有可能让 KK 支持考虑预算？
@rai-studio-bot 也看看
There was a problem hiding this comment.
另外有点好奇，什么样的 case 能在 FF 下不超预算，在 KK 均衡后而且 num_mbs 还经过 VPP round up 之后反而超过预算了
难道是单条数据本身 token 就超过预算了？（FF 确实不会在这里报错，有这种可能）
@rai-studio-bot 如果不是单条就超预算，能构造一个这样的 case 么？
There was a problem hiding this comment.
两个问题一起回（都用 head 11edf3ce 的真实函数验证过）：
1. KK 超预算的具体 case（不是单条超预算，VPP round-up 也无关）
所有样本都不超预算、FF 合法、KK 均衡后超预算的 case 是存在的：
from relax.utils.data.data import get_first_fit_partitions
from relax.utils.data.seqlen_balancing import get_seqlen_balanced_partitions
lengths, cap = [200, 200, 300, 300, 500, 200], 600
ff = get_first_fit_partitions(lengths, cap) # 各包 sums=[600, 600, 500]，全部 <= 600，k=3
kk = get_seqlen_balanced_partitions(lengths, len(ff), equal_size=False)
print([sum(lengths[i] for i in p) for p in kk]) # [700, 500, 500] -> 700 > 600最小版本：lengths=[2, 2, 3, 3, 5, 2], cap=6 → FF 包 sums [6, 6, 5]，KK 包 sums [7, 5, 5]。
机制：num_mbs 是 FF 计数的 DP max，VPP round-up 只会增加包数——我验证了 num_mbs=4（round up 到 2 的倍数）时 KK sums [5, 5, 4, 3]，不超。所以超预算来自 KK 自身：样本数 n > 包数 k 时 KK 必须合并样本，而 KK（largest differencing，见 relax/utils/data/seqlen_balancing.py:47 的 State.merge）只优化各包和的 spread、没有容量输入，合并结果可以越过 FF 的可行装箱。这正是 relax/backends/megatron/data.py:214 的检查与回退存在的原因（回退后 sums [600, 600, 500]）。
单条样本本身超预算也确实会触发（如 [50, 5, 5, 5], cap=40 → FF 单独一包 50 > 40），但这种情况任何打包都无解，回退救不了它——属于序列切分/CP 的范畴，超出打包能做的。
2. 关于「让 KK 支持考虑预算」
可行，但要换一种约束形式。KK 的贪心差分过程没有容量参数，直接改成"带硬容量约束的多路划分"是另一类问题（NP-hard，KK 只是启发式）。两个更现实的落点：
- KK 后修复（改动最小）：保留 KK 结果，只把超限分包里的样本挪到未满的分包（或仅对超限分包做 FF 重排），大多数分包保持 KK 的均衡。需要保证修复后包数不变（
get_data_iterator下游按索引直接取partitions[j]）。 - 换分解方式（slime 的做法，PR 描述里引用的三个链接）：先用带容量的 FF 产出合法包，再把整包分配/均衡到 DP/micro-batch——容量约束天然满足，均衡对象从样本变成包。这是
get_data_iterator的结构级改动，比本 PR 大。
当前"超预算整体回退 FF"作为基线是合理的：正确性优先，均衡损失只影响耗时。如果实际 workload 里回退触发频繁、均衡劣化可测量，KK 后修复是最小的增量改进。
There was a problem hiding this comment.
- 换分解方式（slime 的做法，PR 描述里引用的三个链接）：先用带容量的 FF 产出合法包，再把整包分配/均衡到 DP/micro-batch——容量约束天然满足，均衡对象从样本变成包。这是
get_data_iterator的结构级改动，比本 PR 大。
就是在 FF 后加上 expand_bins_by_splitting 么？@rai-studio-bot
There was a problem hiding this comment.
There was a problem hiding this comment.
接近，但只对了一半——expand_bins_by_splitting 的对应物本 PR 已经有了，slime 方案真正多出来的第三步在 Relax 里映射不进来：
1. FF + 拆包对齐：本 PR 已等价实现。 get_first_fit_partitions 与 slime 的 first_fit_pack 是同一规则（输入顺序 FF、逐包容量检查）；relax/backends/megatron/data.py:219 的拆包循环对应 expand_bins_by_splitting——目标都是包数不足时拆大包，且拆出来的都是原包子集、容量不会增大。差别只在策略：slime 选「token 和最大的多样本包」用 LPT（最长优先进较轻一半）对半拆；本 PR 选「样本数最多」的包每次剥一个样本。这只影响均衡度，不影响容量语义。
2. 关键第三步是「整包按权重分配到 DP」，这一步在 Relax 没有输入。 slime 的 build_dp_schedule 是对整个 step 的全局样本做一次 FF，再把整包按权重（token/FLOPs 和）用 KK 分配到 DP rank（balance_data 时走 get_seqlen_balanced_partitions(weights, dp_size, equal_size=True)）。前提是打包发生在 DP 分配之前、看得到全局样本。而 Relax 的 get_data_iterator 拿到的已经是本 rank 的 local samples（relax/backends/megatron/data.py:959，DP 的样本分配在 rollout 数据分发时已完成，函数内只剩 all-reduce MAX 对齐包数 + dummy 填充），整包分配没有对象可分。要照搬就得把 DP 样本分配挪到打包之后——这才是我说的「结构级改动」的实质。
所以只加 expand_bins_by_splitting 不会带来新东西（对应物已在）。如果想低成本改善回退路径的均衡，倒是可以把 data.py:219 的拆包策略换成 slime 的「按 token 和选包 + LPT 对半拆」：容量语义不变，只改拆哪个包、怎么拆，属于局部改动。
There was a problem hiding this comment.
两点都对：均衡损失不只影响耗时，也影响实际显存峰值；num_mbs 修改前后不变。
显存峰值：动态打包下激活峰值由 token 和最大的那个 micro-batch 决定。修改前 KK 分包没有容量约束，最大包的 token 和可以超过 capacity（KK 只优化各包 spread）；修改后回退保证每个多样本分包 ≤ max_tokens_per_gpu * cp_size（6000 组随机用例中多样本分包 0 次超容量）。例外是单条样本本身超预算（[50, 5, 5, 5], cap=40）：修改前后都是单独一包超容量，打包无解，只有序列切分 / CP 能处理。所以准确的说法是：修改前峰值由 KK 均衡后的最大包决定、可以超 capacity；修改后被 capacity 封顶（单条超预算除外）。
num_mbs：不变。num_mbs 的推导链是 First-Fit 计数（relax/backends/megatron/data.py:1073）→ DP all-reduce MAX（:1078）→ VPP round-up（:1083），本 PR 未改动这条链——get_minimum_num_micro_batch_size 只是改为复用新函数，3000 组随机用例验证计数逐例一致。改动发生在计数之后的分组：KK 分包超容量时回退 First-Fit 并按需拆包（relax/backends/megatron/data.py:214-220），把同样的样本装进同样多的包里。同一批 6000 组用例也验证了回退后包数仍精确等于 num_mbs、dummy offset 正确——这是下游 for j in range(num_mbs) 直接索引的前提。
This branch has not been deployed
3 participants
What
修复普通 RL 动态分包中，KK 重新组合样本后 microbatch 可能超过 token 预算的问题。
Why
开启
--use-dynamic-batch-size后，旧代码先用 First-Fit 计算包数，再经 DP/VPP 对齐，最后由 KK 生成实际分组。问题是 First-Fit 只返回包数，没有保留实际分组；KK 只平衡负载，不限制每包的 token 总量。因此，即使包数正确，重新组合后的 microbatch 仍可能超预算，增加 OOM 风险。
How
参考 slime 默认动态 token 打包的 First-Fit 分组、拆包对齐和 KK 按整包分配到 DP。本 PR 保留 Relax 原有的合法 KK 分组，仅在超预算时回退。