batch size 是怎么被切开的

# verl/workers/engine/utils.py
if use_dynamic_bsz:
    ...
else:
    total_data_size = len(data)
    micro_batch_size_per_gpu = data["micro_batch_size_per_gpu"]
    assert total_data_size % (force_group_size * micro_batch_size_per_gpu) == 0
    micro_batches = tu.chunk_tensordict(
        data, total_data_size // (micro_batch_size_per_gpu * force_group_size)
    )

特性

• 条数固定：每个 micro-batch 恰好 micro_batch_size_per_gpu 条样本，无论它们多长。
• 显存不可控：极端情况一个 micro-batch 全是长样本，token 总数远超预期，OOM。
• DP 间天然对齐：每张卡的 micro-batch 数完全一致，FSDP/Megatron 的集合通信节奏天然同步。
• 失衡浪费算力：短样本 padding 到 batch 里最长那条，长样本拉低吞吐，GPU 利用率低。

入口代码

# verl/workers/engine/utils.py — prepare_micro_batches()
if use_dynamic_bsz:
    assert "max_token_len_per_gpu" in data.keys()
    max_token_len_per_gpu = data["max_token_len_per_gpu"]
    max_token_len = max_token_len_per_gpu * sp_size   # SP 倍率
    micro_batches, batch_idx_list = rearrange_micro_batches(
        data, max_token_len=max_token_len,
        dp_group=dp_group,
        same_micro_num_in_dp=same_micro_num_in_dp,   # 关键
        ...
    )

三步切分流程

步骤 1：算每张卡至少要切几个 micro-batch

total_seqlen     = attention_mask.sum()
num_micro_batches = ceil(total_seqlen / max_token_len)

步骤 2：DP 维度同步 micro-batch 数（这步是 dynamic 模式独有的关键）

if dist.is_initialized() and same_micro_num_in_dp and dp_group is not None:
    num_micro_batches = torch.tensor([num_micro_batches], device=...)
    dist.all_reduce(num_micro_batches, op=dist.ReduceOp.MAX, group=dp_group)
    num_micro_batches = num_micro_batches.cpu().item()

步骤 3：用 Karmarkar-Karp 算法把样本按工作量平均分进 num_micro_batches 个桶

# 工作量估算：模拟 attention FLOPs ≈ 12·h²·L + 2·h·L²
workload = 24576 * seqlen + seqlen ** 2
partitions = karmarkar_karp(workload, k_partitions=num_micro_batches, equal_size=False)

# verl/trainer/ppo/ray_trainer.py — _balance_batch()
global_seqlen_lst = batch.batch["attention_mask"].sum(-1)
workload_lst = calculate_workload(global_seqlen_lst)
dp_size = self._get_dp_size(self.actor_rollout_wg, "actor")

# equal_size=True：每个 rank 拿到等条数样本
global_partition_lst = get_seqlen_balanced_partitions(
    workload_lst, k_partitions=dp_size, equal_size=True
)

# 把 batch 按 KK 给的分组重排，再 chunk 出去
global_idx = torch.tensor([j for partition in global_partition_lst for j in partition])
batch.reorder(global_idx)

这层和 dynamic_bsz 的关键区别

• equal_size=True：每个 rank 必须拿到等条数样本（PPO 后续 mini-batch 数对齐需要），KK 在这个硬约束下做平衡。
• 只重排顺序，不重切样本：原样本不变，只是 chunk 时被分到不同 rank。
• 对 workload 而非 token 数平衡：用 12·h²·L + 2·h·L² 估算 attention FLOPs，让长样本被特别"重视"。
• 每个 mini-batch 跑一次，rollout 完成后、进 forward 前。