让 Packing 加速你的大模型训练

Words:  3367📖   Time:  13min⌛   Category:  AI📌   Tag:   Infra📍

这篇 Blog 介绍一下跟训练效率相关的内容，来自于两篇主要通过 Packing 技术解决 VLM 和 LLM 中 Load Unbalance 的 Paper，都是来自上海交通大学和商汤科技研究院的两项工作

在解读论文之前，先对 Data Packing 进行一个简单的介绍，现在看来它是一种很常见的消除组内 Padding 的方法。如下图所示，左边是我们一般组 batch 的方式，绿色是文本原本的 token，灰色是 Padding 的 token，这样子组 batch 的缺点在于会浪费算力去算灰色的 token（这些token 是原本不需要去计算的）

当我们采用 Packing 的方式去组 batch 时，就不存在这个问题了（如右边）。需要考虑的是怎么样才可以使得组到我们想要的长度（尽量将算力全用到有效的地方）以及会不会给 loss 计算带来问题

OmniBal

介绍的第一篇论文是 Towards Fast Instruction-Tuning for Vision-Language Models via Omniverse Computation Balance，OmniBal 是 Omniverse Balance 的简称。它主要是针对多模态模型训练效率的优化，论文中分别在 DeepSpeed 和 Megatron 上对当时的 Internvl 和 Qwenvl 进行了实验。提出从三个方面来对训练进行优化：data、model 和 memory，接下来我们进行一一介绍

Data Imbalance

经过我开头对 Packing 的介绍，很容易知道一般的组 batch 方式很容易造成 token 计算的浪费。其实除了这个问题，如下图所示我们考虑 DP 层面，由于各个 DP 中 mini-batch 数据长度总和不一样，那么各个 DP 的计算量也就不一样，从而致使可能计算量小的 DP 等待计算量大的 DP 的情况当然你可以将数据都 Padding 到固定长度，但这样的话就会浪费更多 token 的计算 。假如我们能把每个 DP 的数据 Packing 到一样的长度这里是假设计算量和数据总长度大致线性相关时 ，就可以解决这个问题

文中使用了一种称之为 Iterative Sampling Filtering 的算法来进行 Packing，具体流程如下伪代码所示，首先是 sampling 阶段：

def sampling(data_list):
	group = []  # 当前进行 pack 的组，eg: [data1, data2, ...]
	# 存在已经满足条件的组, eg: [[data1, data2, ...], [data4, ...], ...]
	groups = []
	vit_len_sum, total_len_sum = 0, 0
	for data in data_list:
		vit_len, total_len = data
		# 当前的长度满足所设定的最大值 vit_length total_lenth 时，完成该 group 的 packing
		if vit_len_sum + vit_len > vit_length or total_len_sum + total_len > total_lenth:
			groups.append(group)
			vit_len_sum = vit_len
			total_len_sum = total_len
			group = [data]
		else:
			# 否则直接加入该 group
			group.append(data)
			vit_len_sum += vit_len
			total_len_sum += total_len
	groups.append(group)
	return groups

由上面的代码可是得知，组好的 group （除了最后一个）都是在差一条数据的情况就达到组 group 条件之一：group 的 vit 数量之和大于设定 vit 的最大值 vit_length ；group 数据长度之和大于设定的长度最大值 total_lenth 。那么就会存在一个问题，如下图所示黄色部分是阈值长度，group1 和 group2 都差一条数据（灰色部分）达到阈值从而组成了 group，但是他们之间也会存在长度差距

所以，需要筛出离指定阈值差距较大的 group，如下进行 Filtering 阶段：

def filtering(groups):
	filter_groups = []  # 用于存放满足过滤条件的 group
	need_resampling_data = []  # 需要再进行 group 的数据
	for group in groups:
		group_vit_len = sum_vit_len(group)
		group_total_len = sum_total_len(group)
		# total_lenth >= group_total_len > total_length_thresh 满足组好的条件
		if group_vit_len == vit_length or group_total_len >= total_length_thresh:
			filter_groups.append(group)
		else:
			need_resampling_data.extend(group)
		
	return filter_groups, need_resampling_data

对于需要重新 sample 的数据 need_resampling_data，我们可以再次放进 sampling 进行组 group，然后 filtering ，如此循环，体现出了 Iterative Sampling Filtering 中的 Iterative。经过这样的算法，得到的 Packing 结果的总长度可以控制在一定的范围内，使各个 DP 间得到较好的 balance

Model Imbalance & Memory Imbalance

这两个优化是建立在 Data Balance 之上，当 DP 拥有了近似计算量的输入时，需要考虑对于 VLM 怎么去设置 PP 划分是更合适的。其原因是 VLM 是一个异构的架构，单纯的参数或者层均分都可以造成 PP 之间的计算量不均衡的情况

论文使用了一种 profile 每一层的 forward 时间，那么就能得到总的 forward 时间。假如要对模型进行 PP 3 的切分，也就是需要找到中间两层作为切分点，由于有 forward 时间信息，很容易找到能均分 PP 时间的两层，假如是 layer_m 和 layer_n 。另外论文还会找这两层上下层一起作为候选的切分层，因此对于 PP 切分点1 有 [layer_m-1 layer_m layer_m+1] 三层，同理对于 PP 切分点2 也有同样的三层，总的切分办法一共是 9 种。那么怎么衡量这 9 种方法哪种更好呢，文中提出了一下指标：

\(\text{VAR}(\text{fwd_time})\) 是各 PP 的时间与均值时间的差值总和，\(\text{SUM}(\text{comm})\) 是各 PP 之间激活值（应该是每个 PP 中最后一层的输出）通信量的总和。也就是文章考虑了计算和通信两个层面来进行选择

对于 memory 是对于不同长度的输入，来调整 model 需要开启 checkpoint 的层数，本质上是尽量提高显存的利用率，减少重计算的层数，从而提高训练效率

Performance

可以看到，在两种后端的提速表现都非常明显，而且保证了精度没有明显损失的情况

在不同 size 的模型上也进行了速度验证，最快有超过 3 倍的提速！

HBP

HBP 是 Hierarchical Balance Packing 的简称，来自论文 Towards Efficient Supervised Fine-tuning for Long-Context LLM。有了上面的基础，比较容易理解 HBP，实际上这篇论文解决的问题很清晰：训练长文本时，难免出现其他长度的数据或者想和其他长度的数据训练（128k 和 32k 8k 混训），那么如果全都 Packing 到 128k，在开 SP 的情况下会存在 all to all 通信浪费

论文的实验都是基于 DeepSpeed 为后端进行的，整体的 Framework 如下图所示

在 OmniBal 中我们介绍了可以通过调整开启 checkpoint 的层数来减少重计算从而加速训练，那么现在当开启 SP 时，SP 开得越大，空出来的显存越多，多的显存可以用于减少开启 checkpoint 的层数，但是 SP 越大通信就越耗时，所以这俩需要进行均衡取舍的，上图中的 Hierarchical Groups Auto-Selection 就是干这个事

当对于每个 length 确定了该用什么样的 checkpoint 和 SP 后，就可以进行 Hierarchical Packing 了。由上图可知，对数据按照 length 进行分组后，先 Packing 较长的 length，并用短的数据进行尝试填充，使其尽量 Packing 到指定长度，然后再 Packing 短的

接下来详细介绍这两部分以及为了训练稳定做的课程学习和 Loss 修改部分

Hierarchical Groups Auto-Selection

整个流程，首先对每个 length 在开某个 SP 的条件下，尽量去尝试不开 checkpoint 的最大层数，比如对于 32k 开 SP 为 2 对应不开的 checkpoint 的最大层数为 0（即需要全开），继续 profile SP 为 4、8 不开 checkpoint 的层数（尽量不要超过 8 因为节点间通信慢），最后得到 32k 在 SP 为 8 不开 checkpoint 层数为 24 层是 32k 所有配置中最快的，同理我们可以得出 128k、64k 以及 16k、8k 的最快对应的 SP checkpoint 配置，同时我们也可以得出最快的 length 是哪一个，如下图则是对 LLama 8B profile 的配置结果（注意这里的每个 iter 会调整 gradient accumulation step 来使得每个 iter 过的 token 数量一样）

一般来讲，profile 的最大 length 是我们需要训练的最大 length ～ \(l_{max}\)，\(l_{max.sp}\) 为 \(l_{max}\) 最快的 SP 配置，这个 length 一定在我们的训练 config 中。最快的 length（如果不是 max）也是需要加到 config 这里标记为 \(l_{best}\) 和对应 SP \(l_{best.sp}\)，这两个配置是一定有的，另外我们需要考虑长度为 \(l_1=l_{best}/l_{best.sp}\) 以及 \(l_2=l_{max}/l_{max.sp}\) 可以理解为这两个 length 是 max 和 best SP 后每个 rank 跑的 length，对应的 SP 设置为 1（即没有 all to all 通信，一个 rank 是一个 DP）。如果 \(l_2 \lt l_{best}\) 则直接并入 \(l_{best}\)，因此最后最多有四组 Packing length 的设置，最少有三组

Balance Packing

经过了上述的 profile 后，我们可以根据需要训练的 length 对数据进行分组，从 length 最长的组进行 Packing，由 OmniBal 的 ISF 可知，可能存在一些组未到 length_thresh 的标准，这个时候 HBP 的做法则是一次从后面的组取出 sample 进行补充 Packing 直到达到 thresh 或者遍历完，如下为伪代码所示：

# packed_samples 为已经用 ISF Packing 过的 groups
for sample_id_l, sample_l in enumerate(packed_samples):
    llm_num = token_sum(sample_l)
    # 如果 group 长度达到 thresh，即 pack_thresh < llm_num < pack_length 则略过
    if llm_num > packed_thresh:
        continue
    # g_id 越小，里面的数据长度就越小，比如 grouped_samples[0] 都是还未 Packing 的 8k 以下的数据
    for sub_g_id in range(g_id-1, -1, -1):
        for sample_s in grouped_samples[sub_g_id]:
            if llm_num > packed_thresh:
                break
            # 如果该 sample 用于去 Packing 更长数据组了，则需要改正其 g_id
            # 不再加入原长度组的 Packing
            if sample_s['group_id'] == sub_g_id and llm_num + sample_s['token_num'] < packed_length:
                sample_s['group_id'] = g_id
                packed_samples[sample_id_l].append(sample_s)
                llm_num += sample_s['token_num']

最后 Packing 数据的形式会是怎么样呢，如下图所示：

可以看到对于 g_id=2 会有来自短数据的来进行充分 Packing，对于g_id=1 是由两个小包组成一个大包，每个 DP 分到一个小包，如果还需要 SP，则每个 rank 实际跑的数据长度会更小

这里有 2 个细节需要处理：

• 对于每一个层级需要设置不同的 SP group
• 对于每一个层级的 gradient accumulation step 不同

Curriculum Learning & Stable Loss Normalizer

这两个都是为了稳定训练的，前者是通过调整数据训练顺序让短数据先训练一部分，使模型能先适应数据的类型。后者则是为了缓解 Packing 给长短数据带来的影响，比如如果长数据和短数据一起训的时候，由于是长短数据一共的 token 数作为 normalization，那么短数据的 loss 贡献程度必然会被大幅度削减，从而影响较短数据优化。文中使用一个 global batch size 的 token 作为 normalization 实际上使用了一个常数，这样保证梯度不大而且各长度数据 loss 贡献不变

Performance

可以看到加速效果非常明显，这里 balance 是指对于有没有按照包总长度平方进行 sort（这里是对 attention 那的计算量做了一个均衡）

以上便是关于两篇论文的全部内容了

🌸 August 20, 2025