企业官网建设流程全解析

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 Decreasing（FFD）与 Karmarkar-Karp（KK），完整讲解它们的算法原理、复杂度差异、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 与 KK

AReaL 支持两种打包算法，定义于 areal/utils/seqpack.py 的PACKING_ALGORITHMS注册表中（PACKING_ALGORITHM_FFD = "ffd"、PACKING_ALGORITHM_KK = "kk"）：

FFD：贪心装箱启发式

FFD 的实现位于 seqpack.py，核心逻辑分四步：

1. 将序列按长度降序排序（np.argsort(-values)）；
2. 若当前分组数少于min_groups，直接新建分组；
3. 否则在所有能容纳当前序列的分组中，选择当前总长度最小的那个放入；
4. 若都放不下，新建一个分组。

实现中通过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）。完整的字段说明如下：

YAML 配置

在实验配置文件中，通过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" (默认), "kk"

Python API

也可以直接在代码中构造或修改MicroBatchSpec：

from areal.api.cli_args import MicroBatchSpec # 使用 KK 算法 mb_spec = MicroBatchSpec( max_tokens_per_mb=8192, n_mbs=4, packing_algorithm="kk", ) # 或者更新一个现有的 spec mb_spec_kk = MicroBatchSpec.new(existing_spec, packing_algorithm="kk")

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, capacity=int(1e12), min_groups=world_size)——故意将容量设为极大值、只按world_size做纯均衡切分，最后把第 i 组轨迹分配给 rank i。这正是 KK 算法"忽略容量、专注最小化 spread"特性的典型生产场景：保证每个 DP rank 分到的 token 总量尽可能一致。

3. 算法分发器

get_allocate_fn（seqpack.py）是两种算法的统一分发入口，对未知算法名抛出包含可用选项的ValueError，保证配置错误能尽早暴露。

何时使用 KK，何时 FFD 就足够

KK 的推荐场景

• 序列长度变化极大的大规模 RL 训练（如 RLHF、开放式生成的 PPO）：KK 显著缩小负载最重和最轻的 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_pct（spread 改善百分比）等指标——这正是文档所述"双峰分布 + 高并行度"场景的实验化验证。

何时 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 的均衡性不差于 FFD；
• tests/test_seqpack.py 对既有ffd_allocate做回归测试，防止重构破坏行为。

小结与选型建议

对于 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