使用Hugging Face交错数据集时遇UnboundLocalError问题排查

错误信息 #

Exception has occurred: UnboundLocalError
local variable 'batch_idx' referenced before assignment
  File "/lfs/skampere1/0/brando9/beyond-scale-language-data-diversity/src/training/utils.py", line 254, in _test_train_dataset_setup_for_main_code
    print(f'{len(next(iter(batch))["input_ids"])=}')
  File "/lfs/skampere1/0/brando9/beyond-scale-language-data-diversity/src/training/utils.py", line 263, in <module>
    _test_train_dataset_setup_for_main_code()
UnboundLocalError: local variable 'batch_idx' referenced before assignment

触发场景 #

该错误在处理Hugging Face流式（streaming）数据集时触发，涉及数据集加载、分词、文本分组等预处理步骤，相关代码片段如下：

raw_train_datasets = load_dataset(path[0], name[0], data_files=data_files[0], streaming=streaming, split=split[0]).with_format("torch")
get_data_from_hf_dataset(raw_train_datasets, streaming=streaming, batch_size=batch_size) 
remove_columns = get_column_names(raw_train_datasets)
tokenize_function = lambda examples: tokenizer(examples["text"])
tokenized_train_datasets = raw_train_datasets.map(tokenize_function, batched=True, remove_columns=remove_columns)
block_size: int = tokenizer.model_max_length
_group_texts = lambda examples : group_texts(examples, block_size)
lm_train_dataset = tokenized_train_datasets.map(_group_texts, batched=True)
batch = get_data_from_hf_dataset(lm_train_dataset, streaming=streaming, batch_size=batch_size)
print(f'{len(next(iter(batch))["input_ids"])=}')

错误原因 #

这是Hugging Face datasets库的内部bug：当流式数据集被提前设置with_format("torch")后，后续执行map等转换操作会导致数据集内部状态异常，迭代数据时触发batch_idx变量未初始化的错误。

代码中提前对原始数据集应用with_format("torch")，破坏了流式数据集在转换过程中的内部迭代逻辑，最终在尝试获取批次数据时触发该异常。

解决方案 #

延迟设置with_format("torch")到所有预处理操作完成之后，避免格式转换干扰流式数据集的内部处理流程：

修改后的代码示例 #

# 加载数据集时不设置torch格式
raw_train_datasets = load_dataset(path[0], name[0], data_files=data_files[0], streaming=streaming, split=split[0])
get_data_from_hf_dataset(raw_train_datasets, streaming=streaming, batch_size=batch_size) 
remove_columns = get_column_names(raw_train_datasets)

# 执行分词、文本分组等预处理
tokenize_function = lambda examples: tokenizer(examples["text"])
tokenized_train_datasets = raw_train_datasets.map(tokenize_function, batched=True, remove_columns=remove_columns)
block_size: int = tokenizer.model_max_length
_group_texts = lambda examples : group_texts(examples, block_size)
lm_train_dataset = tokenized_train_datasets.map(_group_texts, batched=True)

# 所有预处理完成后，再设置torch格式
lm_train_dataset = lm_train_dataset.with_format("torch")

# 后续正常获取批次数据
batch = get_data_from_hf_dataset(lm_train_dataset, streaming=streaming, batch_size=batch_size)
print(f'{len(next(iter(batch))["input_ids"])=}')

关键注意点 #

• 流式数据集的格式转换（with_format）必须放在所有map、interleave_datasets等转换操作之后
• 避免在原始流式数据集上直接绑定格式，保持其原始状态直到所有预处理完成

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