fix(megatron): repack oversized KK microbatches

[DrRyanHuang]

What

修复普通 RL 动态分包中，KK 重新组合样本后 microbatch 可能超过 token 预算的问题。

Why

开启 --use-dynamic-batch-size 后，旧代码先用 First-Fit 计算包数，再经 DP/VPP 对齐，最后由 KK 生成实际分组。

问题是 First-Fit 只返回包数，没有保留实际分组；KK 只平衡负载，不限制每包的 token 总量。因此，即使包数正确，重新组合后的 microbatch 仍可能超预算，增加 OOM 风险。

How

MegatronTrainRayActor.train_actor()
  ├─ get_data_iterator()
  │    First-Fit 计算包数 → DP MAX / VPP 对齐包数
  │      → _get_seqlen_partitions_with_dummy_padding()
  │           KK 分组 → 容量检查
  │             合法：保留 KK
  │             超限：First-Fit 重打包 → 必要时拆包对齐
  └─ model.train() → train_one_step()
       → Megatron 前反向调度 → forward_step()
       → get_batch() → DataIterator.get_next()

• 增加返回实际样本索引分组的 First-Fit 函数，原计数函数复用其结果。
• 将容量传到实际分组函数，KK 超预算时回退 First-Fit。
• 若包数不足 DP/VPP 对齐要求，只拆已有包直到满足目标包数。

参考 slime 默认动态 token 打包的 First-Fit 分组、拆包对齐和 KK 按整包分配到 DP。本 PR 保留 Relax 原有的合法 KK 分组，仅在超预算时回退。

[rai-studio-bot]

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 pre-commit run --all-files before submitting.

请说明问题、改动和验证方式，保持 PR 聚焦，并在提交前运行 pre-commit run --all-files。

English contribution guide · 中文贡献指南

CI commands / CI 指令

Command / 指令	Usage / 用途
/rerun	Retry failed CI / 重跑失败的 CI
/rerun <target>	Rerun a workflow or check / 重跑指定 workflow 或检查
/cancel <workflow>	Cancel an entire workflow / 取消整个 workflow
/help	Show commands and targets / 查看指令和 target
/review	Request a code review / 请求代码 review

Put one command on the first line of a new PR comment. Rerun/cancel require PR authorship or repository write access.

在新 PR 评论的首行写一条指令。PR 作者或有仓库写权限的贡献者可以重跑、取消 CI。

CI usage and targets · CI 用法与 target

[rai-studio-bot]

Nyanpasu 审查看板

审查状态： 💬 已完成 · 有补充意见

审查版本： 11edf3c

目标分支： main

审查完成，结论为 Comment：打包回退机制本身经验证正确（计数等价 3000 组、回退正确性 6000 组随机用例），无阻塞问题；提出 1 个 P2 测试缺口（F1）与 1 个 P3 非行级跟进建议（F2，SFT 路径同样缺少 KK 容量上限）。

审查阶段	进度	范围与结果
常规审查	✅ 已完成	已审查全部 2 个变更文件的回退实现、全部调用方与既有测试覆盖；随机模拟验证 First-Fit 计数重构前后逐例一致、回退后覆盖完整且多样本分包不超容量；pre-commit 通过；CI 当前无失败。
深度审查	✅ 已完成	必要性审计（生产代码与测试）由父任务完成：评估了容量感知 KK、只重打包超限分包、删除计数包装函数、合并两处拆包循环等替代方案并给出保留/补充结论；独立设计参照因变更范围窄而跳过（理由已记录），无未解决的深度问题。

审查发现

待处理

编号	严重性	问题	状态	规则来源
F1		为容量回退路径补一个纯函数回归测试	🚧 未解决	—
F2		SFT 预打包路径同样缺少 KK 容量上限（非行级跟进建议）	🚧 未解决	—

已解决或已取代

编号	严重性	问题	状态	规则来源
暂无已解决或已取代的记录。

提交范围 · 接收 2 · 建议移出 0 · 待确认 0

接收 2 个文件 · 建议移出 0 个文件 · 待确认 0 个文件。移出与待确认部分暂停深审，不代表审查通过。

文件	结论	仓库维护必要性	依据	替代去向或方案
relax/backends/megatron/data.py relax/utils/data/data.py	接收	Both files implement the fix itself: get_first_fit_partitions now returns index partitions and _get_seqlen_partitions_with_dummy_padding falls back to it when KK partitions exceed the token budget. get_data_iterator is the production entry point for RL dynamic packing with existing callers in the Megatron training loop, and get_minimum_num_micro_batch_size is reused by relax/backends/megatron/actor.py:1461.	PR description: use-dynamic-batch-size First-Fit count without retained grouping, KK balancing without capacity bound; verified callers at relax/backends/megatron/data.py:1073 and relax/backends/megatron/actor.py:1461.	No smaller placement satisfies the contract: the fallback needs actual index partitions, so a partition-returning first-fit helper is required. Keeping it beside the pre-existing get_minimum_num_micro_batch_size in relax/utils/data/data.py is the minimal change; moving it into seqlen_balancing.py would not reduce concepts.

精简审查与验证依据

审查范围	进度	结论
生产代码	✅ 已完成	两个生产文件的机制均有明确契约与调用方支撑；评估过的替代方案均因损失语义或收益不足而保留。
测试	✅ 已完成	无测试文件变更；现有相关测试家族保护不同边界，均保留。本次修复的核心回退路径无测试覆盖，已作为 F1 提出（补充而非替换）。

生产代码的必要性与替代方案

范围	必须保留的契约	更简单的方案	结论	依据与限制
relax/utils/data/data.py：get_first_fit_partitions 及 get_minimum_num_micro_batch_size 包装函数	First-Fit 计数语义（按输入顺序装箱、超长样本独占一包）需保持；现有调用方 relax/backends/megatron/actor.py:1461 与 relax/backends/megatron/data.py:1073 依赖该计数对齐 DP/VPP 包数。	删除包装函数，两个调用点直接写 len(get_first_fit_partitions(...))。	保留	包装函数保留具名的计数概念，仅 2 个调用方；内联不减少概念反而让包数语义散落。3000 组随机用例验证重构前后计数逐例一致，行为保持。
relax/backends/megatron/data.py：_get_seqlen_partitions_with_dummy_padding 的容量检查、First-Fit 回退与拆包循环，及 get_data_iterator 的 capacity 变量提取	动态分包后每个 micro-batch 的 token 总量不超过 max_tokens_per_gpu * cp_size，且最终包数必须精确等于 num_mbs（下游 for j in range(num_mbs) 直接索引 partitions）。	把 KK 改成容量感知（类似 slime 按整包分配到 DP），或仅对超限分包单独重打包而非整体回退；与既有 _split_preference_bins_to_count 合并拆包循环。	保留	容量回退是本 PR 的目的，不可删除；整体回退 + 拆包远比容量感知 KK 简单，且与上游 slime 的 First-Fit 分组/拆包对齐一致。6000 组随机用例确认回退后包数恰等于 num_partitions、多样本分包不超容量。两处拆包循环的选择策略与顺序语义不同（preference 需保持相邻插入），合并需参数化两个变体，收益有限。

测试的必要性与替代方案

范围	必须保留的契约	更简单的方案	结论	依据与限制
现有 tests/backends/megatron/test_data_vpp.py（仅 use_dynamic_batch_size=False）与 tests/utils/data/test_seqlen_balancing.py（仅 KK 分组）	非动态路径的 rollout mini 边界行为，以及 KK 分组的平衡性与覆盖性。	合并为一个共享测试文件。	保留	两个家族分别保护入口行为与纯算法边界，合并无收益。动态路径与新的回退分支完全无覆盖（get_data_iterator 动态分支、get_first_fit_partitions 均无测试），该缺口由 F1 承载；建议的回归测试为纯函数级补充，不引入重复维护。

^{Powered by Nyanpasu with glm-5.3-flash max, please check the suggestions carefully.}

[rai-studio-bot]

审查结论

本 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)

• P3 · 非行级：问题位于本 PR 未改动的文件，无法挂到 diff 行 · SFT 预打包路径同样缺少 KK 容量上限

  relax/backends/megatron/actor.py:1462（以及 :1294 的 DP 对齐重分包）同样先以 First-Fit 计数、再用 get_seqlen_balanced_partitions 分组，且没有容量检查，与本 PR 修复前的 RL 路径模式相同。若 SFT 的 pack 长度可能严重不均，会有同样的超预算风险。本 PR 明确限定 RL 路径，此点留作后续跟进即可。

^{Powered by Nyanpasu with glm-5.3-flash max, please check the suggestions carefully.}

[rai-studio-bot]

P2 · 为容量回退路径补一个纯函数回归测试

回退逻辑本身经验证是正确的：用真实的 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 风险。