Huggingface Dataset.map设batched=True时进度条变红未完成的原因

复现代码 #

model_checkpoint = "distilgpt2"

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)

def tokenize_function(examples):
    return tokenizer(examples["text"])

from datasets import load_dataset
datasets = load_dataset('wikitext', 'wikitext-2-raw-v1')
tokenized_datasets = datasets.map(tokenize_function, batched=False, num_proc=4, remove_columns=["text"])

问题描述 #

当设置batched=False时，进度条显示绿色表示执行成功，但设置batched=True时，进度条显示红色且无法到达100%。请问这是否意味着我的map函数执行失败，还是存在其他原因？

回答 #

进度条变红且无法完成确实说明map执行出了问题，核心原因和代码逻辑、运行环境的适配有关：

• 当batched=False时，函数每次处理单个样本，tokenizer能直接处理单条字符串，多进程模式下也不会有资源冲突，因此顺利完成；
• 当batched=True时，map会把多个样本打包成列表传给tokenize_function，虽然tokenizer本身支持批量输入，但结合num_proc=4多进程时，很容易触发以下问题：
  • 内存占用过高：批量处理时单进程需要加载更多数据，多进程叠加后超出内存上限，导致进程崩溃；
  • 多进程资源冲突：部分环境下多进程处理批量数据时，会出现tokenizer资源无法共享的问题，引发进程挂起。

解决建议 #

• 先去掉num_proc=4，用单进程测试batched=True，如果能正常运行，就说明是多进程和批量处理的兼容性问题；
• 如果单进程也卡，手动设置batch_size参数（比如batch_size=32），减小每次处理的样本量，降低内存压力；
• 确认tokenize_function的返回结构：批量处理时，tokenizer返回的每个字段（比如input_ids）都应该是二维列表，确保和map的要求匹配。

内容的提问来源于stack exchange，提问作者DiveIntoML