dataset 路由列注入导致混合 mmap/内存 Arrow 表在多进程 map 中出现序列化与内存放大

Open
  • Is something wrong?
#10,543 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
42/100
Issue type
Bug
Clarity
Mostly clear
Activity status
Active
Tech stack
python
Domain
backend, data-engineering

Research direction

Start with _inject_dataset_routing_tag() in swift/dataset/loader.py and compare the current add_column() behavior with the batched map() reproduction in the issue. Reproduce the reported pickle size increase using the supplied synthetic script and its listed dependency versions. Done means routing values are preserved while multi-process shard serialization avoids the reported memory amplification; the issue does not specify a target implementation.

Written by the indexing model from the issue text.

Description

bug
Checklist / 检查清单
  • I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues，确认这是一个新的 bug report。
Bug Description / Bug 描述
问题描述

_inject_dataset_routing_tag() 使用 Dataset.add_column() 为非 streaming 数据集注入常量 dataset 路由列。对于磁盘数据集，原有列是 mmap 表，新增列则是内存表。

在多数据块、shuffle 和逻辑分片的组合下，新增路由列会导致后续多进程 map 的任务序列化显著放大。普通 SFT 未启用多教师时也会注入该列。

最小合成实验中，同一个 32 行逻辑 shard 的序列化大小：

  • 原始磁盘表：3.7 KB；
  • add_column() 注入路由列后：3.30 MB；
  • 分批 map 注入并落盘后：1.5 KB。

希望保留路由列语义，改进其注入方式，避免将磁盘数据集变成容易产生序列化放大的混合表。

受影响代码

截至 2026-11-10，main 和 v4.6.0 仍使用以下实现；历史受影响环境为 ms-swift 4.5.3。

def _inject_dataset_routing_tag(dataset: DATASET_TYPE, ds_name: str) -> DATASET_TYPE:
    if isinstance(dataset, HfIterableDataset):
        return dataset.map(lambda example: {**example, 'dataset': ds_name})
    if 'dataset' in dataset.column_names:
        dataset = dataset.remove_columns('dataset')
    return dataset.add_column('dataset', [ds_name] * len(dataset))

源码：swift/dataset/loader.py。

How to Reproduce / 如何复现
最小复现

CPU-only，仅使用合成数据。脚本通过 multiprocess.reduction.ForkingPickler 比较多进程任务使用的数据集 shard 序列化大小，不需要模型或 GPU。临时 Arrow 文件退出时自动清理。

依赖：datasets、pyarrow、multiprocess。

import os
import sys
import tempfile

import datasets
import multiprocess
import pyarrow
from datasets import Dataset, concatenate_datasets, load_from_disk
from multiprocess.reduction import ForkingPickler


def main():
    datasets.disable_progress_bars()
    print(
        "versions:",
        "Python", sys.version.split()[0],
        "datasets", datasets.__version__,
        "pyarrow", pyarrow.__version__,
        "multiprocess", multiprocess.__version__,
    )

    n_chunks = 16
    rows_per_chunk = 32
    tag = "synthetic://" + "x" * 185

    with tempfile.TemporaryDirectory(prefix="swift-routing-") as root:
        pieces = []
        for i in range(n_chunks):
            path = os.path.join(root, str(i))
            Dataset.from_dict({
                "text": [f"row-{i}-{j}" for j in range(rows_per_chunk)]
            }).save_to_disk(path)
            pieces.append(load_from_disk(path, keep_in_memory=False))

        disk = concatenate_datasets(pieces)

        # 模拟 Swift 对非 streaming 数据集注入路由标签。
        mixed = disk.add_column("dataset", [tag] * len(disk))

        # 对照组：通过分批 map 将路由标签物化到磁盘。
        mapped = disk.map(
            lambda batch: {"dataset": [tag] * len(batch["text"])},
            batched=True,
            batch_size=rows_per_chunk,
            writer_batch_size=rows_per_chunk,
            keep_in_memory=False,
            cache_file_name=os.path.join(root, "routing.arrow"),
            load_from_cache_file=False,
        )

        assert mixed["text"] == mapped["text"]
        assert mixed["dataset"] == mapped["dataset"]

        for name, ds in [
            ("disk_only", disk),
            ("add_column", mixed),
            ("batched_map_disk", mapped),
        ]:
            shard = ds.shuffle(seed=42).shard(num_shards=n_chunks, index=0)
            print(
                name,
                "logical_rows", len(shard),
                "backing_rows", shard._data.num_rows,
                "pickle_bytes", len(ForkingPickler.dumps(shard)),
            )


if __name__ == "__main__":
    main()
实测结果

2026-11-10 运行环境：Python 3.10.15、datasets 3.6.0、pyarrow 20.0.0、multiprocess 0.70.14。

disk_only logical_rows 32 backing_rows 512 pickle_bytes 3717
add_column logical_rows 32 backing_rows 512 pickle_bytes 3304107
batched_map_disk logical_rows 32 backing_rows 512 pickle_bytes 1497

路由列总共只有 512 行，字符串逻辑数据量约 0.10 MB，但一个 32 行 shard 序列化后达到 3.30 MB。两种带路由列的实现具有相同的数据行和路由值。

原因是逻辑分片仍引用完整的底层表，新增内存列的切片共享父缓冲区，序列化混合表时重复携带这些缓冲区。放大发生在后续任务准备/传输阶段，磁盘表本身仍通过 mmap 路径重建。

较大规模实验记录

2026-10-17 环境：ms-swift 4.5.3、datasets 4.8.4、pyarrow 25.0.1、multiprocess 0.70.19。

96 个底层数据块、9,600 行、196 字符路由标签，shuffle 后取 100 行逻辑 shard：

数据布局 shard 序列化大小（字节）
仅磁盘列 9,493
add_column() 新增路由列 368,690,531
带路由列的数据集保存并重新加载 1,449

971,079 行、96 worker 的离线预处理对照：

方案 cgroup 匿名内存峰值 已进入编码的 worker 数 停止条件
原始路由列注入 262.71 GiB 1 256 GiB 保护阈值触发主动终止
分批 map 注入并落盘 93.46 GiB 96 900 秒观察窗口结束

以上为历史实验记录；两组均未完成全量 map，展示的是启动期内存差异。

与现有 issue / PR 的关系
  • ms-swift #10121 修复的是路由列注入触发的多 rank 缓存竞争，仍保留 add_column()；这里报告的是新增内存列的序列化放大。
  • ms-swift #10483 讨论线程过度订阅和 template/tokenizer 的重复 pickle。本复现不包含 template/tokenizer，放大对象是数据集 shard。
  • ms-swift #10533 报告多模态长度预处理随进程数增加而变慢，尚未定位路由列或混合表序列化问题，可作为相关讨论关联。
  • datasets #7662 讨论新增列后多进程 map OOM，维护者建议通过分批 map 生成新列。
  • datasets #8611 的历史测试版本（25e99170168dc5c482048c32282e043e372a67a3）仅压紧不含磁盘 cache files 的内存 shard，没有覆盖本例的混合表布局。
Additional Information / 补充信息

No response

Dominant language
Python
Stars
16.1k
Forks
1.8k
Avg merge
1d 17h
Merged PRs (30d)
70

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from modelscope/ms-swift

All issues in modelscope/ms-swift

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.