seqlen_balancing.py — 序列长度负载均衡

文件路径: verl/utils/seqlen_balancing.py

文件概述

seqlen_balancing.py 实现了基于 Karmarkar-Karp 差分算法的序列长度负载均衡。在分布式训练中，如果不同 GPU 分到的序列长度差异很大，快的 GPU 会等慢的，造成浪费。这个模块将序列尽量均匀分配。

背景知识

负载均衡问题: 将 N 个不等长序列分配到 K 个 GPU 上，使各 GPU 处理的总 token 数尽量相等。这是一个 NP-hard 的数字划分问题。

Karmarkar-Karp 算法 是该问题的著名近似算法，通过不断将最大的两个数合并（差分），得到近似最优的划分。

核心函数详解

1. 分配算法

def get_seqlen_balanced_partitions(lengths, num_partitions, equal_size=False):
    """
    将序列按长度均衡地分配到 num_partitions 个分区。

    Args:
        lengths: 每个序列的长度列表
        num_partitions: 分区数（等于 GPU 数）
        equal_size: 是否要求每个分区的序列数相等

    Returns:
        List[List[int]]: 每个分区包含的序列索引
    """

2. 工作流程

输入: 序列长度 [100, 200, 50, 300, 150, 80]
分配到 2 个 GPU:

GPU 0: [300, 100, 80]  → 总长 480
GPU 1: [200, 150, 50]  → 总长 400

比简单轮询的 [300+50+150=500, 200+100+80=380] 更均衡

与其他模块的关系

• 被 DataLoader 在分发批次到各 GPU 时调用
• 配合 tensordict_utils.py 重排 TensorDict

小结

seqlen_balancing.py 用经典的负载均衡算法解决了分布式训练中序列长度不均的问题，减少了 GPU 空闲等待时间。