聊聊HuggingFace如何处理大模型下海量数据集

翻译自：
Big data? 🤗 Datasets to the rescue!

如今，使用大GB的数据集并不罕见，特别是从头开始预训练像BERT或GPT-2这样的Tranformer模型。在这样的情况下，甚至连加载数据都可能是一个挑战。例如，用于预训练GPT-2的WebText语料库包含超过800万份文档和40GB的文本——将其加载到电脑的RAM中，可能会使其炸掉。
幸运的是，HuggingFace DataSet 数据集的设计已经克服了这些限制，它通过将数据集视为内存映射文件(Memory-mapped file)来解决内存管理问题，并通过流(Streaming)处理语料库中的条目来解决硬盘驱动器限制。
我们将使用一个巨大的825 GB语料库(称为Pile)来探索HuggingFace数据集的这些特性。

什么是Pile?

Pile是一个英语文本语料库，由EleutherAI创建，用于训练大规模语言模型。它包括各种各样的数据集，涵盖科学文章、GitHub代码库和过滤后的web文本。训练语料库以14GB块的形式提供，你还可以下载几个单独的组件。
从PubMed Abstracts数据集开始，这是PubMed上1500万份生物医学出版物的摘要语料库。数据集是JSON行格式，并使用zstandard库压缩，所以首先我们需要安装它:

!pip install zstandard

接下来，可以使用HuggingFace提供的数据集下载方式来加载：

from datasets import load_dataset

# This takes a few minutes to run, so go grab a tea or coffee while you wait :)
data_files = "https://the-eye.eu/public/AI/pile_preliminary_components/PUBMED_title_abstracts_2019_baseline.jsonl.zst"
pubmed_dataset = load_dataset("json", data_files=data_files, split="train")
pubmed_dataset

Dataset({
  features: ['meta', 'text'],
  num_rows: 15518009
})

可以看到，这里有15518009行、2列的数据。可以查看下输出的数据集内容的第一个示例：

print pubmed_dataset[0];

# output:
{'meta': {'pmid': 11409574, 'language': 'eng'},
 'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection.\nTo determine the prevalence of hypoxaemia in children aged under 5 years suffering acute lower respiratory infections (ALRI), the risk factors for hypoxaemia in children under 5 years of age with ALRI, and the association of hypoxaemia with an increased risk of dying in children of the same age ...'}

这看起来像是一篇医学文章的摘要。 现在让我们看看我们使用了多少 RAM 来加载数据集！

memory mapping

在 Python 中测量内存使用情况的一个简单方法是使用 psutil 库，可以使用 pip 安装该库，如下所示：

!pip install psutil

它提供了一个Process类，允许我们检查当前进程的内存使用情况，如下所示：

import psutil

# Process.memory_info is expressed in bytes, so convert to megabytes
print(f"RAM used: {psutil.Process().memory_info().rss / (1024 * 1024):.2f} MB")

# output:
RAM used: 5678.33 MB

这里的 rss 属性指的是驻留集大小，它是进程在 RAM 中占用的内存部分。 此测量还包括 Python 解释器和我们加载的库使用的内存，因此用于加载数据集的实际内存量要小一些。 为了进行比较，我们使用 dataset_size 属性查看数据集在磁盘上的大小。 由于结果像以前一样以字节表示，因此我们需要手动将其转换为千兆字节：

print(f"Number of files in dataset : {pubmed_dataset.dataset_size}")
size_gb = pubmed_dataset.dataset_size / (1024**3)
print(f"Dataset size (cache file) : {size_gb:.2f} GB")

# output
Number of files in dataset : 20979437051
Dataset size (cache file) : 19.54 GB

很好——尽管它有近 20 GB 大，但我们能够用更少的 RAM 加载和访问数据集！
如果你熟悉 Pandas，这个结果可能会让你感到惊讶，因为 Wes Kinney 著名的经验法则是，你通常需要的 RAM 是数据集大小的 5 到 10 倍。 那么HuggingFace数据集是如何解决这个内存管理问题的呢？ HuggingFace Datasets 将每个数据集视为内存映射文件，它提供 RAM 和文件系统存储之间的映射，允许库访问和操作数据集的元素，而无需将其完全加载到内存中。
内存映射文件还可以在多个进程之间共享，这使得 Dataset.map() 等方法可以并行化，而无需移动或复制数据集。 在底层，这些功能都是由 Apache Arrow 内存格式和 pyarrow 库实现的，这使得数据加载和处理速度快如闪电。 为了查看实际情况，让我们通过迭代 PubMed Abstracts 数据集中的所有元素来运行一些速度测试：

import timeit

code_snippet = """batch_size = 1000

for idx in range(0, len(pubmed_dataset), batch_size):
    _ = pubmed_dataset[idx:idx + batch_size]
"""

time = timeit.timeit(stmt=code_snippet, number=1, globals=globals())
print(
    f"Iterated over {len(pubmed_dataset)} examples (about {size_gb:.1f} GB) in "
    f"{time:.1f}s, i.e. {size_gb/time:.3f} GB/s"
)

# output:
'Iterated over 15518009 examples (about 19.5 GB) in 64.2s, i.e. 0.304 GB/s'

这里我们使用Python的timeit模块来测量code_snippet所花费的执行时间。 你通常能够以十分之几 GB/秒到几 GB/秒的速度迭代数据集。 这对于绝大多数应用程序来说都非常有效，但有时你必须使用太大而无法存储在笔记本电脑硬盘上的数据集。 例如，如果我们尝试下载整个 Pile，我们将需要 825 GB 的可用磁盘空间！ 为了处理这些情况，Hugging Face Datasets 提供了流式传输功能，允许我们动态下载和访问元素，而无需下载整个数据集。

Streaming Datasets

要启用数据集流式传输，你只需将Streaming=True参数传递给load_dataset()函数。 例如，让我们再次加载 PubMed Abstracts 数据集，但采用流模式：

pubmed_dataset_streamed = load_dataset(
    "json", data_files=data_files, split="train", streaming=True
)

Streaming=True 返回的对象不是我们在本章其他地方遇到的熟悉的 Dataset，而是 IterableDataset。 顾名思义，要访问 IterableDataset 的元素，我们需要对其进行迭代。 我们可以访问流数据集的第一个元素，如下所示：

next(iter(pubmed_dataset_streamed))

# output
{'meta': {'pmid': 11409574, 'language': 'eng'},
 'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection.\nTo determine the prevalence of hypoxaemia in children aged under 5 years suffering acute lower respiratory infections (ALRI), the risk factors for hypoxaemia in children under 5 years of age with ALRI, and the association of hypoxaemia with an increased risk of dying in children of the same age ...'}

可以使用IterableDataset.map()即时处理流数据集中的元素，如果你需要对输入进行标记，这在训练期间非常有用。

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
tokenized_dataset = pubmed_dataset_streamed.map(lambda x: tokenizer(x["text"]))
next(iter(tokenized_dataset))

# output
{'input_ids': [101, 4958, 5178, 4328, 6779, ...], 'attention_mask': [1, 1, 1, 1, 1, ...]}

还可以使用IterableDataset.shuffle()对流式数据集进行混洗，但与 Dataset.shuffle() 不同，它仅对预定义的 buffer_size 中的元素进行混洗：

shuffled_dataset = pubmed_dataset_streamed.shuffle(buffer_size=10_000, seed=42)
next(iter(shuffled_dataset))

# output
{'meta': {'pmid': 11410799, 'language': 'eng'},
 'text': 'Randomized study of dose or schedule modification of granulocyte colony-stimulating factor in platinum-based chemotherapy for elderly patients with lung cancer ...'}

在此示例中，我们从缓冲区中的前 10,000 个示例中随机选择了一个示例。 一旦访问了一个示例，它在缓冲区中的位置就会被语料库中的下一个示例填充（即上述情况中的第 10,001 个示例）。 还可以使用 IterableDataset.take() 和IterableDataset.skip()函数从流式数据集中选择元素，其作用方式与Dataset.select()类似。 例如，要选择 PubMed Abstracts 数据集中的前 5 个示例，我们可以执行以下操作：

dataset_head = pubmed_dataset_streamed.take(5)
list(dataset_head)

# output
[{'meta': {'pmid': 11409574, 'language': 'eng'},
  'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection ...'},
 {'meta': {'pmid': 11409575, 'language': 'eng'},
  'text': 'Clinical signs of hypoxaemia in children with acute lower respiratory infection: indicators of oxygen therapy ...'},
 {'meta': {'pmid': 11409576, 'language': 'eng'},
  'text': "Hypoxaemia in children with severe pneumonia in Papua New Guinea ..."},
 {'meta': {'pmid': 11409577, 'language': 'eng'},
  'text': 'Oxygen concentrators and cylinders ...'},
 {'meta': {'pmid': 11409578, 'language': 'eng'},
  'text': 'Oxygen supply in rural africa: a personal experience ...'}]

同样，可以使用 IterableDataset.skip() 函数从打乱的数据集中创建训练和验证分割，如下所示：

# Skip the first 1,000 examples and include the rest in the training set
train_dataset = shuffled_dataset.skip(1000)
# Take the first 1,000 examples for the validation set
validation_dataset = shuffled_dataset.take(1000)

让我们通过一个常见的应用程序来完成对数据集流的探索：将多个数据集组合在一起以创建单个语料库。 🤗 Datasets 提供了 interleave_datasets() 函数，它将 IterableDataset 对象列表转换为单个 IterableDataset，其中新数据集的元素是通过源示例之间的交替获得的。 当尝试合并大型数据集时，此函数特别有用，因此我们以流式传输 Pile 的 FreeLaw 子集为例，该子集是来自美国法院的 51 GB 法律意见数据集：

law_dataset_streamed = load_dataset(
    "json",
    data_files="https://the-eye.eu/public/AI/pile_preliminary_components/FreeLaw_Opinions.jsonl.zst",
    split="train",
    streaming=True,
)
next(iter(law_dataset_streamed))

{'meta': {'case_ID': '110921.json',
  'case_jurisdiction': 'scotus.tar.gz',
  'date_created': '2010-04-28T17:12:49Z'},
  'text': '\n461 U.S. 238 (1983)\nOLIM ET AL.\nv.\nWAKINEKONA\nNo. 81-1581.\nSupreme Court of United States.\nArgued January 19, 1983.\nDecided April 26, 1983.\nCERTIORARI TO THE UNITED STATES COURT OF APPEALS FOR THE NINTH CIRCUIT\n*239 Michael A. Lilly, First Deputy Attorney General of Hawaii, argued the cause for petitioners. With him on the brief was James H. Dannenberg, Deputy Attorney General...'}

这个数据集足够大，足以给大多数笔记本电脑的 RAM 带来压力，但我们已经能够毫不费力地加载和访问它！ 现在让我们将 FreeLaw 和 PubMed Abstracts 数据集中的示例与 interleave_datasets() 函数结合起来：

from itertools import islice
from datasets import interleave_datasets

combined_dataset = interleave_datasets([pubmed_dataset_streamed, law_dataset_streamed])
list(islice(combined_dataset, 2))

[{'meta': {'pmid': 11409574, 'language': 'eng'},
  'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection ...'},
 {'meta': {'case_ID': '110921.json',
   'case_jurisdiction': 'scotus.tar.gz',
   'date_created': '2010-04-28T17:12:49Z'},
  'text': '\n461 U.S. 238 (1983)\nOLIM ET AL.\nv.\nWAKINEKONA\nNo. 81-1581.\nSupreme Court of United States.\nArgued January 19, 1983.\nDecided April 26, 1983.\nCERTIORARI TO THE UNITED STATES COURT OF APPEALS FOR THE NINTH CIRCUIT\n*239 Michael A. Lilly, First Deputy Attorney General of Hawaii, argued the cause for petitioners. With him on the brief was James H. Dannenberg, Deputy Attorney General...'}]

在这里，我们使用 Python 的 itertools 模块中的 islice() 函数从组合数据集中选择前两个示例，我们可以看到它们与两个源数据集中的第一个示例相匹配。

最后，如果想流式传输整个 825 GB 的 Pile，可以按如下方式获取所有准备好的文件：

base_url = "https://the-eye.eu/public/AI/pile/"
data_files = {
    "train": [base_url + "train/" + f"{idx:02d}.jsonl.zst" for idx in range(30)],
    "validation": base_url + "val.jsonl.zst",
"test": base_url + "test.jsonl.zst",
}
pile_dataset = load_dataset("json", data_files=data_files, streaming=True)
next(iter(pile_dataset["train"]))

# output
{'meta': {'pile_set_name': 'Pile-CC'},
 'text': 'It is done, and submitted. You can play “Survival of the Tastiest” on Android, and on the web...'}

总结

总结来看，主要是通过内存映射与流处理来实现的大数据集加载，这也是业界比较常用的方案。