AReaL 序列打包算法FFD 与 KK面向 RL 训练的微批次负载均衡详解【免费下载链接】AReaLThe RL Bridge for LLM-based Agent Applications. Made Simple Flexible.项目地址: https://gitcode.com/GitHub_Trending/are/AReaL导读本文聚焦 AReaL 训练流程中控制微批次micro-batch分配的两大序列打包算法——First Fit DecreasingFFD与 Karmarkar-KarpKK完整讲解它们的算法原理、复杂度差异、YAML 与 Python API 配置方式并结合仓库源码与测试用例揭示其在数据并行DP负载均衡中的实际作用。读完本文你将掌握如何针对不同序列长度分布与并行规模选择打包算法并理解 KK 在变长序列 RL 训练中显著缩小 DP rank 间负载差距spread的底层原因。为什么要关心序列打包在 AReaL 这类面向 LLM Agent 应用的强化学习训练框架中rollout 阶段生成的序列长度天然是变长的不同 prompt 的回复长度差异巨大尤其在 RLHF、PPO 这类开放式生成场景下尤为明显。训练侧需要把这些变长序列切分成多个微批次micro-batch喂给模型而如何分组直接决定了每个 DP rank 承担多少 token 计算量分组均衡时各 rank 的负载接近同步屏障barrier处的等待时间很短训练吞吐量高分组失衡时负载最重的 rank 成为瓶颈其余 rank 被迫在 barrier 处空转等待。因此序列打包不仅是装箱问题更是吞吐量优化问题。AReaL 通过MicroBatchSpec中的packing_algorithm字段将打包策略暴露为可配置选项让用户在不同训练场景下自由切换。支持的两种算法FFD 与 KKAReaL 支持两种打包算法定义于 areal/utils/seqpack.py 的PACKING_ALGORITHMS注册表中PACKING_ALGORITHM_FFD ffd、PACKING_ALGORITHM_KK kk算法Key描述复杂度均衡质量First Fit Decreasing (FFD)ffd贪心装箱启发式算法。按长度降序对序列排序并将每个序列分配到第一个还有剩余容量的桶中。O(n log n)良好 (Good)Karmarkar-Karp (KK)kk最大差分法 (Largest Differencing Method)。使用最大堆迭代合并两个最不平衡的部分分区产生接近最优的均衡效果。O(n log n · k)极佳 (Excellent)FFD贪心装箱启发式FFD 的实现位于 seqpack.py核心逻辑分四步将序列按长度降序排序np.argsort(-values)若当前分组数少于min_groups直接新建分组否则在所有能容纳当前序列的分组中选择当前总长度最小的那个放入若都放不下新建一个分组。实现中通过bisect维护一个按分组总长度有序的列表group_values保证放入最轻的分组操作在 O(log n) 内完成。同时ffd_allocate还保证输出分组数不小于min_groups且能被n_groups_divisor整除不满足时自动向上补齐min_groups后重试。这与 MicroBatchSpec 的字段语义 完全一致——n_mbs是最小微批次数n_mbs_divisor约束微批次数必须被其整除常用于流水线并行。KK最大差分法KK 算法Karmarkar-Karp Largest Differencing Method的实现位于 seqpack.py核心思路是迭代合并两个最不平衡的部分分区每个候选状态_KKState持有 k 个集合_KKSet并维护最大堆__lt__按spread max_sum - min_sum降序定义即 spread 最大的状态最先被弹出每次从堆中弹出 spread 最大的两个状态将它们交叉配对合并merge时把self的最大集合与other的最小集合配对反之亦然从而最小化合并后的 spread反复合并直到堆中只剩一个状态即为最终分区。算法参考了 R.E. Korf 的《Multi-Way Number Partitioning》IJCAI 2009。值得注意的实现细节kk_allocate分组数由max(min_groups, ceil(total / capacity))决定并向上取整到n_groups_divisor的倍数但不能超过序列总数这与 veRL 的rearrange_micro_batches计算num_micro_batches ceildiv(total_seqlen, max_token_len)的做法一致当capacity被设为极大值如int(1e12)时算法忽略容量约束、纯粹追求分组均衡——这一用法在轨迹重分配trajectory redistribution中被实际采用安全网机制若 KK 产生的某个分组仍超过容量会记录 warning 并自动回退到 FFD保证任何配置下都不会产生超容量的微批次。打包质量度量为了让 KK 与 FFD 的差异可量化seqpack.py 提供了_compute_packing_metrics可输出spread最大-最小负载差、imbalance_ratio、std_dev、cv变异系数、utilization、wasted_tokens等一整套指标用于对比两种算法的均衡效果。配置方式打包算法由MicroBatchSpec中的packing_algorithm字段控制该 dataclass 定义于 areal/api/cli_args.py默认值为ffd可选值仅ffd与kk__post_init__会校验非法值并抛出ValueError。完整的字段说明如下字段默认值说明n_mbs1微批次数当设置了max_tokens_per_mb时表示最小微批次数granularity1每个微批次的粒度相邻序列按此大小分组max_tokens_per_mbNone每个微批次 forward pass 允许的最大 token 数设置后n_mbs退化为下限n_mbs_divisor1最终微批次数会被调整为该值的倍数流水线并行常用packing_algorithmffd打包算法ffd默认或kkYAML 配置在实验配置文件中通过actor.mb_spec下配置参见 examples/countdown/train_config.yaml 的结构以及 tests/grpo/config.yaml 中的实际用法actor: mb_spec: max_tokens_per_mb: 8192 n_mbs: 4 n_mbs_divisor: 1 packing_algorithm: kk # 选项: ffd (默认), kkPython API也可以直接在代码中构造或修改MicroBatchSpecfrom areal.api.cli_args import MicroBatchSpec # 使用 KK 算法 mb_spec MicroBatchSpec( max_tokens_per_mb8192, n_mbs4, packing_algorithmkk, ) # 或者更新一个现有的 spec mb_spec_kk MicroBatchSpec.new(existing_spec, packing_algorithmkk)MicroBatchSpec.new是一个保留 Omegaconf 兼容性的工厂方法会基于现有 spec 的全部字段n_mbs、granularity、max_tokens_per_mb、n_mbs_divisor、packing_algorithm重建一个新实例只覆盖传入的 kwargs——在需要从配置读取后局部修改的场景非常实用。打包算法在训练流程中的实际调用链序列打包并非孤立功能而是深度嵌入 AReaL 的训练数据管线与 rollout 轨迹重分配流程1. 训练侧微批次分配areal/utils/data.py 中的allocate_balanced_mbs是训练侧的核心入口它断言max_tokens_per_mb必须被设置随后通过get_allocate_fn(mb_spec.packing_algorithm)分发到ffd_allocate或kk_allocate把序列长度列表lens按容量、最小分组数、除数约束打包成微批次返回每个微批次包含的序列下标。其配套函数allocate_balanced_mbs_synced还通过all_gather_object在进程组内对齐各 rank 的微批次数确保 DP 拓扑下所有 rank 的微批次结构一致。2. 推理侧轨迹重分配areal/infra/dist_rollout.py 的redistribute_trajectories展示了另一种用法把全体 rank 的 rollout 轨迹按attention_mask求和得到真实序列长度然后调用allocate_fn(seqlens, capacityint(1e12), min_groupsworld_size)——故意将容量设为极大值、只按world_size做纯均衡切分最后把第 i 组轨迹分配给 rank i。这正是 KK 算法忽略容量、专注最小化 spread特性的典型生产场景保证每个 DP rank 分到的 token 总量尽可能一致。3. 算法分发器get_allocate_fnseqpack.py是两种算法的统一分发入口对未知算法名抛出包含可用选项的ValueError保证配置错误能尽早暴露。何时使用 KK何时 FFD 就足够KK 的推荐场景序列长度变化极大的大规模 RL 训练如 RLHF、开放式生成的 PPOKK 显著缩小负载最重和最轻的 DP rank 之间的差距spread双峰序列分布 (bimodal sequence distributions)极短与极长序列混合时贪心打包容易把长序列挤到一起造成失衡而 KK 的最大差分策略能更合理地抵消长短序列高 DP 并行度≥4 个 rank此时即使很小的负载不平衡也会因同步屏障导致明显的空闲等待。仓库的分布式对比脚本 tests/torchrun/run_kk_vs_ffd.py 给出了一个可直接验证上述场景的模拟实验它用generate_bimodal_seqlens构造短序列 50~200 token、长序列 800~2048 token、长序列占比 30%的双峰分布在 4 个 rank 上分别用真实实现ffd_allocate/kk_allocate重分配轨迹并记录ffd_spread、kk_spread、kk_wins、improvement_pctspread 改善百分比等指标——这正是文档所述双峰分布 高并行度场景的实验化验证。何时 FFD 就足够了均匀或接近均匀的序列长度相比均衡度更关注打包开销的小规模实验FFD 为 O(n log n)更快对延迟敏感的推理流水线FFD 速度略快。测试覆盖与配置校验AReaL 为打包功能提供了充分的测试保障tests/test_kk_allocate.py 系统性覆盖kk_allocate的容量/最小分组/除数约束、容量超限报错、equal_size模式、get_allocate_fn分发正确性ffd返回ffd_allocate、kk返回kk_allocate、未知算法抛异常、MicroBatchSpec字段校验默认值为ffd、非法值被拒绝并包含 KK 与 FFD 的随机化对比测试tests/test_kk_allocate.py断言 KK 的均衡性不差于 FFDtests/test_seqpack.py 对既有ffd_allocate做回归测试防止重构破坏行为。小结与选型建议判断维度选 FFD选 KK序列长度分布均匀 / 接近均匀高度可变、双峰分布并行规模小规模实验DP 4大规模 RL、高 DP 并行度≥4首要目标打包开销低、延迟敏感负载均衡最小化 spread复杂度O(n log n)O(n log n · k)对于 AReaL 上典型的 RL 训练rollout 序列长度差异大、DP 并行度高默认优先尝试packing_algorithm: kk而均匀长度或对延迟敏感的推理场景保持默认的ffd即可。两种算法均可通过MicroBatchSpec一行配置切换且 KK 自带容量超限回退 FFD 的安全网切换成本极低。【免费下载链接】AReaLThe RL Bridge for LLM-based Agent Applications. Made Simple Flexible.项目地址: https://gitcode.com/GitHub_Trending/are/AReaL创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 AReaL 序列打包算法（FFD 与 KK）：面向 RL 训练的微批次负载均衡详解
AReaL 序列打包算法FFD 与 KK面向 RL 训练的微批次负载均衡详解【免费下载链接】AReaLThe RL Bridge for LLM-based Agent Applications. Made Simple Flexible.项目地址: https://gitcode.com/GitHub_Trending/are/AReaL导读本文聚焦 AReaL 训练流程中控制微批次micro-batch分配的两大序列打包算法——First Fit DecreasingFFD与 Karmarkar-KarpKK完整讲解它们的算法原理、复杂度差异、YAML 与 Python API 配置方式并结合仓库源码与测试用例揭示其在数据并行DP负载均衡中的实际作用。读完本文你将掌握如何针对不同序列长度分布与并行规模选择打包算法并理解 KK 在变长序列 RL 训练中显著缩小 DP rank 间负载差距spread的底层原因。为什么要关心序列打包在 AReaL 这类面向 LLM Agent 应用的强化学习训练框架中rollout 阶段生成的序列长度天然是变长的不同 prompt 的回复长度差异巨大尤其在 RLHF、PPO 这类开放式生成场景下尤为明显。训练侧需要把这些变长序列切分成多个微批次micro-batch喂给模型而如何分组直接决定了每个 DP rank 承担多少 token 计算量分组均衡时各 rank 的负载接近同步屏障barrier处的等待时间很短训练吞吐量高分组失衡时负载最重的 rank 成为瓶颈其余 rank 被迫在 barrier 处空转等待。因此序列打包不仅是装箱问题更是吞吐量优化问题。AReaL 通过MicroBatchSpec中的packing_algorithm字段将打包策略暴露为可配置选项让用户在不同训练场景下自由切换。支持的两种算法FFD 与 KKAReaL 支持两种打包算法定义于 areal/utils/seqpack.py 的PACKING_ALGORITHMS注册表中PACKING_ALGORITHM_FFD ffd、PACKING_ALGORITHM_KK kk算法Key描述复杂度均衡质量First Fit Decreasing (FFD)ffd贪心装箱启发式算法。按长度降序对序列排序并将每个序列分配到第一个还有剩余容量的桶中。O(n log n)良好 (Good)Karmarkar-Karp (KK)kk最大差分法 (Largest Differencing Method)。使用最大堆迭代合并两个最不平衡的部分分区产生接近最优的均衡效果。O(n log n · k)极佳 (Excellent)FFD贪心装箱启发式FFD 的实现位于 seqpack.py核心逻辑分四步将序列按长度降序排序np.argsort(-values)若当前分组数少于min_groups直接新建分组否则在所有能容纳当前序列的分组中选择当前总长度最小的那个放入若都放不下新建一个分组。实现中通过bisect维护一个按分组总长度有序的列表group_values保证放入最轻的分组操作在 O(log n) 内完成。同时ffd_allocate还保证输出分组数不小于min_groups且能被n_groups_divisor整除不满足时自动向上补齐min_groups后重试。这与 MicroBatchSpec 的字段语义 完全一致——n_mbs是最小微批次数n_mbs_divisor约束微批次数必须被其整除常用于流水线并行。KK最大差分法KK 算法Karmarkar-Karp Largest Differencing Method的实现位于 seqpack.py核心思路是迭代合并两个最不平衡的部分分区每个候选状态_KKState持有 k 个集合_KKSet并维护最大堆__lt__按spread max_sum - min_sum降序定义即 spread 最大的状态最先被弹出每次从堆中弹出 spread 最大的两个状态将它们交叉配对合并merge时把self的最大集合与other的最小集合配对反之亦然从而最小化合并后的 spread反复合并直到堆中只剩一个状态即为最终分区。算法参考了 R.E. Korf 的《Multi-Way Number Partitioning》IJCAI 2009。值得注意的实现细节kk_allocate分组数由max(min_groups, ceil(total / capacity))决定并向上取整到n_groups_divisor的倍数但不能超过序列总数这与 veRL 的rearrange_micro_batches计算num_micro_batches ceildiv(total_seqlen, max_token_len)的做法一致当capacity被设为极大值如int(1e12)时算法忽略容量约束、纯粹追求分组均衡——这一用法在轨迹重分配trajectory redistribution中被实际采用安全网机制若 KK 产生的某个分组仍超过容量会记录 warning 并自动回退到 FFD保证任何配置下都不会产生超容量的微批次。打包质量度量为了让 KK 与 FFD 的差异可量化seqpack.py 提供了_compute_packing_metrics可输出spread最大-最小负载差、imbalance_ratio、std_dev、cv变异系数、utilization、wasted_tokens等一整套指标用于对比两种算法的均衡效果。配置方式打包算法由MicroBatchSpec中的packing_algorithm字段控制该 dataclass 定义于 areal/api/cli_args.py默认值为ffd可选值仅ffd与kk__post_init__会校验非法值并抛出ValueError。完整的字段说明如下字段默认值说明n_mbs1微批次数当设置了max_tokens_per_mb时表示最小微批次数granularity1每个微批次的粒度相邻序列按此大小分组max_tokens_per_mbNone每个微批次 forward pass 允许的最大 token 数设置后n_mbs退化为下限n_mbs_divisor1最终微批次数会被调整为该值的倍数流水线并行常用packing_algorithmffd打包算法ffd默认或kkYAML 配置在实验配置文件中通过actor.mb_spec下配置参见 examples/countdown/train_config.yaml 的结构以及 tests/grpo/config.yaml 中的实际用法actor: mb_spec: max_tokens_per_mb: 8192 n_mbs: 4 n_mbs_divisor: 1 packing_algorithm: kk # 选项: ffd (默认), kkPython API也可以直接在代码中构造或修改MicroBatchSpecfrom areal.api.cli_args import MicroBatchSpec # 使用 KK 算法 mb_spec MicroBatchSpec( max_tokens_per_mb8192, n_mbs4, packing_algorithmkk, ) # 或者更新一个现有的 spec mb_spec_kk MicroBatchSpec.new(existing_spec, packing_algorithmkk)MicroBatchSpec.new是一个保留 Omegaconf 兼容性的工厂方法会基于现有 spec 的全部字段n_mbs、granularity、max_tokens_per_mb、n_mbs_divisor、packing_algorithm重建一个新实例只覆盖传入的 kwargs——在需要从配置读取后局部修改的场景非常实用。打包算法在训练流程中的实际调用链序列打包并非孤立功能而是深度嵌入 AReaL 的训练数据管线与 rollout 轨迹重分配流程1. 训练侧微批次分配areal/utils/data.py 中的allocate_balanced_mbs是训练侧的核心入口它断言max_tokens_per_mb必须被设置随后通过get_allocate_fn(mb_spec.packing_algorithm)分发到ffd_allocate或kk_allocate把序列长度列表lens按容量、最小分组数、除数约束打包成微批次返回每个微批次包含的序列下标。其配套函数allocate_balanced_mbs_synced还通过all_gather_object在进程组内对齐各 rank 的微批次数确保 DP 拓扑下所有 rank 的微批次结构一致。2. 推理侧轨迹重分配areal/infra/dist_rollout.py 的redistribute_trajectories展示了另一种用法把全体 rank 的 rollout 轨迹按attention_mask求和得到真实序列长度然后调用allocate_fn(seqlens, capacityint(1e12), min_groupsworld_size)——故意将容量设为极大值、只按world_size做纯均衡切分最后把第 i 组轨迹分配给 rank i。这正是 KK 算法忽略容量、专注最小化 spread特性的典型生产场景保证每个 DP rank 分到的 token 总量尽可能一致。3. 算法分发器get_allocate_fnseqpack.py是两种算法的统一分发入口对未知算法名抛出包含可用选项的ValueError保证配置错误能尽早暴露。何时使用 KK何时 FFD 就足够KK 的推荐场景序列长度变化极大的大规模 RL 训练如 RLHF、开放式生成的 PPOKK 显著缩小负载最重和最轻的 DP rank 之间的差距spread双峰序列分布 (bimodal sequence distributions)极短与极长序列混合时贪心打包容易把长序列挤到一起造成失衡而 KK 的最大差分策略能更合理地抵消长短序列高 DP 并行度≥4 个 rank此时即使很小的负载不平衡也会因同步屏障导致明显的空闲等待。仓库的分布式对比脚本 tests/torchrun/run_kk_vs_ffd.py 给出了一个可直接验证上述场景的模拟实验它用generate_bimodal_seqlens构造短序列 50~200 token、长序列 800~2048 token、长序列占比 30%的双峰分布在 4 个 rank 上分别用真实实现ffd_allocate/kk_allocate重分配轨迹并记录ffd_spread、kk_spread、kk_wins、improvement_pctspread 改善百分比等指标——这正是文档所述双峰分布 高并行度场景的实验化验证。何时 FFD 就足够了均匀或接近均匀的序列长度相比均衡度更关注打包开销的小规模实验FFD 为 O(n log n)更快对延迟敏感的推理流水线FFD 速度略快。测试覆盖与配置校验AReaL 为打包功能提供了充分的测试保障tests/test_kk_allocate.py 系统性覆盖kk_allocate的容量/最小分组/除数约束、容量超限报错、equal_size模式、get_allocate_fn分发正确性ffd返回ffd_allocate、kk返回kk_allocate、未知算法抛异常、MicroBatchSpec字段校验默认值为ffd、非法值被拒绝并包含 KK 与 FFD 的随机化对比测试tests/test_kk_allocate.py断言 KK 的均衡性不差于 FFDtests/test_seqpack.py 对既有ffd_allocate做回归测试防止重构破坏行为。小结与选型建议判断维度选 FFD选 KK序列长度分布均匀 / 接近均匀高度可变、双峰分布并行规模小规模实验DP 4大规模 RL、高 DP 并行度≥4首要目标打包开销低、延迟敏感负载均衡最小化 spread复杂度O(n log n)O(n log n · k)对于 AReaL 上典型的 RL 训练rollout 序列长度差异大、DP 并行度高默认优先尝试packing_algorithm: kk而均匀长度或对延迟敏感的推理场景保持默认的ffd即可。两种算法均可通过MicroBatchSpec一行配置切换且 KK 自带容量超限回退 FFD 的安全网切换成本极低。【免费下载链接】AReaLThe RL Bridge for LLM-based Agent Applications. Made Simple Flexible.项目地址: https://gitcode.com/GitHub_Trending/are/AReaL创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考