高效的多模态数据管道

你已经准备好了一切——数据、模型、强大的 GPU 配置。你点击“运行”，然后……等待。再等一会儿。你的 GPU 几乎没有出什么力，而你的钱包却在每小时变轻。

听起来很熟悉？我们也经历过。在对我们的 nanoVLM 项目进行了一些调查后，我们发现真正的罪魁祸首不是我们的模型或硬件，而是我们的数据管道效率极其低下。

以下是我们的发现

1. GPU 空闲：我们的模型实际上在等待数据送达
2. 填充地狱：每个批次都塞满了无用的填充令牌，这些令牌对训练毫无贡献

在这篇文章中，我们将分五个阶段构建一个高效的管道。在每个阶段，我们都会在前一步的基础上进行增删，并评论哪些地方做得对，哪些地方做得不好。

目录：

• 第 0 阶段：准备工作
• 第 1 阶段：可视化数据集
• 第 2 阶段：朴素填充
• 第 3 阶段：约束填充
• 第 4 阶段：使用背包算法进行更智能的打包
• 第 5 阶段：将背包算法应用于多模态数据
• 结论

【第 0 阶段】准备工作

为了更容易地跟进数据准备任务，我们创建了一个专门关注数据管道的独立仓库。我们希望这比在与 nanoVLM 仓库集成后阅读代码更容易理解。此外，这对于启动其他数据管道也可能很有用！

仓库：https://github.com/ariG23498/mmdp

要跟上进度，您只需克隆该仓库即可。它包含了最终的数据准备任务，但其设计旨在展示过程中的每一步。

$ git clone https://github.com/ariG23498/mmdp.git

【第 1 阶段】可视化数据集

在优化任何东西之前，我们需要了解我们正在处理什么。我们的多模态数据集包含图像、文本提示和响应。

$ uv run 01_check_dataset.py

熟悉你的训练数据对于成功至关重要。前面的脚本每次运行时都会显示一个随机样本；你可能需要将代码片段复制到笔记本中并多次运行，以便对数据有一个感觉。

【第 2 阶段】朴素填充

我们第一次训练尝试采用了显而易见（且非常常见）的方法

• 对所有内容进行分词
• 找到每个批次中最长的序列
• 将其他所有内容填充以匹配长度

$ uv run 02_naive_pad_dataloader.py

结果令人痛苦。看看这个可视化图

看到那些灰色的部分了吗？那就是填充。那是 GPU 在你支付计算时间的同时，却在处理绝对空无一物的东西。我们大约浪费了 60% 的批次在空令牌上。

【第 3 阶段】约束填充

我们的下一步很简单。设置一个全局最大长度并遵守它。如果一个样本太长，我们就直接丢弃它。

你可能已经注意到，批次现在少了一个样本。这是由于过滤过程。这有所帮助，但我们仍然将所有内容填充到相同的固定长度，而不管实际内容如何。比以前好，但仍然很浪费。

【第 4 阶段】：使用背包算法进行更智能的打包

现在我们准备完全重新思考批处理。填充是敌人，我们需要一个策略来最小化它，同时最大化每个批次中可以容纳的数据量。进入背包问题，这是计算机科学中的一个经典问题，非常适合这个场景。

想象一下你正在为一次徒步旅行打包背包。它的承重量有限，而你希望尽可能多地塞进有用的物品。在我们的情况下

• 背包是一个具有最大令牌限制（max_length）的训练批次。
• 每个物品是一个序列（一个分词后的提示-响应对），其重量是令牌的数量。
• 我们的目标是在不超过令牌限制的情况下，将尽可能多的序列打包到批次中，从而最大限度地减少浪费的空间。

为了测试这个想法，我们从一个玩具数据集开始：只是一个从 1 到 25 的数字列表，每个数字代表一个序列长度。这让我们可以在不涉及图像和文本复杂性的情况下进行实验。

切换到可迭代数据集

大多数 PyTorch 数据集是 map-style 的（你用 dataset[i] 访问它们）。但对于动态批处理，我们需要更灵活的东西。因此，我们通过子类化 torch.utils.data.IterableDataset 构建了一个 iterable-style 数据集。这让我们能够动态生成批次，并处理像在多个工作进程之间分片数据这样的技巧。

def _get_data_range(self):
    worker_info = get_worker_info()
    if worker_info is None:  # single worker, return the entire dataset
        return self.start, self.end
    else:  # multiple workers, split the data load
        per_worker = int(
            math.ceil((self.end - self.start) / worker_info.num_workers)
        )
        worker_id = worker_info.id
        iter_start = self.start + worker_id * per_worker
        iter_end = min(iter_start + per_worker, self.end)
        return iter_start, iter_end

生产者-消费者模式的魔力

打包序列可能会很慢，尤其是在我们进行排序或打乱时。为了保持流畅，我们使用生产者-消费者模式，借助 Python 队列

def _producer(self, data_iter, queue, stop_signal):
    if self.strategy == "greedy":
        for pack in self._greedy_packing(data_iter):
            queue.put(pack)
    elif self.strategy == "binpack":
        while True:
            buffer = list(itertools.islice(data_iter, self.buffer_size))
            if not buffer:
                break
            knapsacks = self._bin_packing(buffer)
            for pack in knapsacks:
                queue.put(pack)
    queue.put(stop_signal)

生产者线程打包批次并将其放入队列，而主线程则根据需要从中取出。这种重叠使得管道能够顺畅流动。

贪婪打包

首先，我们尝试一种简单的贪婪打包策略

def _greedy_packing(self, iterator):
    pack, pack_sum = [], 0
    for item in iterator:
        if item > self.max_length:
            continue
        if pack_sum + item <= self.max_length:
            pack.append(item)
            pack_sum += item
        else:
            yield pack
            pack = [item]
            pack_sum = item
    if pack:
        yield pack

这种方法按顺序遍历数据，将项目添加到一个包中直到满，然后开始一个新的包。它速度快但不完美。以下是批次的样子

=== Strategy: GREEDY ===
[tensor([1]), tensor([2]), tensor([3]), tensor([4]), tensor([5]), tensor([6]), tensor([7]), tensor([8]), tensor([9]), tensor([10]), tensor([11]), tensor([12]), tensor([13])]
[tensor([14]), tensor([15]), tensor([16]), tensor([17]), tensor([18]), tensor([19])]
[tensor([20]), tensor([21]), tensor([22]), tensor([23])]
[tensor([24])]

注意到后面的批次变得稀疏了吗？我们留下了空隙。

使用装箱算法实现更紧密的拟合

让我们尝试一种更智能的方法：装箱算法（具体来说是首次适应递减法）

def _bin_packing(self, buffer: List[int]):
    buffer = sorted(buffer, reverse=True)
    knapsacks = []
    for item in buffer:
        for pack in knapsacks:
            if sum(pack) + item <= self.max_length:
                pack.append(item)
                break
        else:
            knapsacks.append([item])

这种方法按长度对序列进行排序（最长的在前），并尝试将每个序列装入第一个有空间的包中。如果都装不下，就开启一个新的包。结果如何？

=== Strategy: BINPACK ===
[tensor([24]), tensor([23]), tensor([22]), tensor([21]), tensor([10])]
[tensor([20]), tensor([19]), tensor([18]), tensor([17]), tensor([16]), tensor([9]), tensor([1])]
[tensor([15]), tensor([14]), tensor([13]), tensor([12]), tensor([11]), tensor([8]), tensor([7]), tensor([6]), tensor([5]), tensor([4]), tensor([3]), tensor([2])]

这些批次要紧密得多，浪费的空间更少。这就像玩俄罗斯方块一样，把数据块紧密地拼接在一起。

【第 5 阶段】将背包算法应用于多模态数据

现在是动真格的时候了，将背包打包算法应用到我们的多模态数据集中。

我们又回到了图像、提示和响应，我们需要在遵守令牌限制和图像预算的同时高效地打包它们。图像预算的设置是为了平衡每个样本的图像数量。我们希望避免一个 GPU 需要处理比另一个 GPU 多得多的图像的情况。

我们新的 ConstantLengthDataset 类负责繁重的工作。以下是它与第 4 阶段相比的工作方式

ConstantLengthDataset 包揽了一切

• 读取样本（图像和文本）。
• 过滤掉太长或图像太多的样本。
• 使用贪婪背包策略将样本打包成批次，平衡令牌数和图像数。
• 将最终批次填充到固定长度，但填充量比以前少得多。

这是结果

看那个！灰色（填充）部分非常少，批次里充满了有用的数据。这就像打包行李箱一样，装得那么好，你不用坐在上面就能拉上拉链。

这张图乍一看可能不直观，但让我们把它和约束填充的图并排比较一下。

在这里你会注意到，背包算法中的样本分布更均匀。我们也不会因为过滤而导致批次中的样本数量减少。

结论

最初只是一个简单的“为什么训练这么慢？”的调查，最终导致我们彻底重新思考了处理多模态数据的方式。

用于数据管道的平衡背包策略来自 NVIDIA 的论文 Eagle 2: Building Post-Training Data Strategies from Scratch for Frontier Vision-Language Models。

关键经验

• 将所有内容填充到最长序列是一个不错的初步方法（但很浪费）
• 将批处理视为一个打包问题
• 考虑所有约束（文本长度、图像内存等）
• 首先用玩具数据测试以验证你的方法

想深入了解吗？请查看

• 我们的 nanoVLM 博客文章
• nanoVLM GitHub
• 这个管道的代码

祝你训练愉快（愿你的 GPU 保持忙碌）！