[LLMのFinetuningのための基礎] datasetsの基礎を理解する

はじめに

今回は [LLMのFinetuningのための基礎] datasetsの基礎を理解する と題しまして、datasetsに関して学習したことをまとめました。

学習教材としてはHuggingFaceのdatasetsのページを使用しています。

学習のアウトプットなので不正確・不十分な説明の可能性もあるため参考にされる場合は十分にご注意ください。

目次

1. datasetsには2種類ある

2. datasetsの重要コンポーネント・関数

3. Tips

4. 終わりに

datasetsには2種類ある

datasetsは目的と用途に合わせて「Dataset」と「IterableDataset」の2種類が用意されています。

IterableDatasetは名前の通りイテラブルにデータを読み込むことが可能です。

これは数十、数百GBといったRAM内に収まりきらないデータセットを扱う際に有効なアプローチになります。

Dataset

Datasetのメリット

Datasetのメリットをいくつか挙げます。

1. ランダムアクセスが可能
2. 全体量を可視化することが可能
3. データを簡単にシャッフル可能

1. ランダムアクセスが可能

Datasetはランダムアクセスが可能です。

したがって特定のインデックスを指定してデータを取り出すことができます。

# 例：
# datasetの先頭を確認する

first_data = dataset[0]

print(first_data)

これはIterableDatasetとは違ってデータセットを全て読み込んでいるからできることだと解釈できます。

より厳密には最初にデータセット全体をRAMに記録しているのではなく、Apache Arrowのようなメモリマッピング技術によって高速化されています。
この詳細について知りたい方はabout_arrowのページかApache Arrowのページをご覧ください。

2. 全体量を可視化することが可能

上記で述べた「これはIterableDatasetとは違ってデータセットを全て読み込んでいるからできることだと解釈できます。」という話に繋がりますが、全体量を確認することも簡単に可能です。

# 例：
# datasetの全体量を確認する

dataset_length = len(dataset)

print(dataset_length)

3. データを簡単にシャッフル可能

こちらも同様です。

データ全体を保持しているためシャッフルも簡単に行えます。

# 例：
# datasetをシャッフルする

shuffle_dataset = dataset.shuffle(seed=42)

Datasetのデメリット

Datasetは非常に柔軟で扱いやすいdatasetsコンポーネントですが、データを全て保持するためRAM内に収まりきらない様な非常に巨大なデータの場合は適していません。

IterableDataset

IterableDatasetのメリット

IterableDatasetのメリットをいくつか挙げます。

1. 非常に巨大なデータセットを扱うことが可能
2. イテラブルに読み込むことが可能

1. 非常に巨大なデータセットを扱うことが可能

Datasetのデメリットで述べたようにRAM内に収まりきらない巨大なデータセットを扱うタイミングはGenAIを訓練する際には非常によくあります。

IterableDatasetではデータを1つずつ取得、処理が可能なため最低限のRAM使用量で済みます。

Tips

IterableDatasetだとしても例えば、for文で最後まで処理を行った場合はRAMがパンパンにならないのでしょうか？

IterableDatasetではデータを読込み → 処理 を行った後にデータはガベージコレクションに送られます。

これによってメモリ領域が解放されるため、for文で最後まで処理を行った場合でもRAMはパンパンにはなりません。

2. イテラブルに読み込むことが可能

これは文字通りのメリットなのですが、例えばリアルタイムのデータを処理したい時や動的に変わってしまうデータを扱うときにDatasetとは違ってその都度読み込むことができるのが大きな利点です。

IterableDatasetのデメリット

IterableDatasetにもいくつかデメリットがあります。

これはDatasetとほぼ真逆なのですが、データセット全体を読み込んでいないため「全体サイズの確認」や「データのシャッフル」といった操作が困難になります。

DatasetとIterableDatasetの相互変換

DatasetとIterableDatasetは相互に変換が可能です。

（どのようなタイミングでそれらが必要になるのかは私の経験不足のため例を挙げられないことご理解ください。）

DatasetからIterableDatasetに変換

こちらはDatasetに搭載されている to_iterable_dataset() で可能です。

iterable_dataset = dataset.to_iterable_dataset()

IterableDatasetからDatasetに変換

こちらも何らかの方法（dor文など）でIterableDatasetの中身を全て取り出したりキャッシュ化したのちにDatasetに変換可能です。

しかし、基本的にはおすすめしません。

理由としては元々IterableDatasetのメリット、使用タイミングとしてはRAM内に収まりきらないほどの大規模データセットを扱う時なので、Dataset化しようとするとメモリーエラーになる可能性が非常に高いからです。

datasetsの重要コンポーネント・関数

datasetsのうち特によく使う・目にするコンポーネント・関数をいくつかご紹介します。

load_dataset関数・Dataset.from_xxx関数

load_dataset関数・Dataset.from_xxx関数（例えば、from_json()など）は任意のDataset・IterableDatasetオブジェクトを作成します。

load_dataset()を使用することで、JSONやCSV、テキストや画像ファイルと言った様々なファイルをDatasetもしくはItelableDataset形式にロードすることができます。

詳しい扱いはload_datasetのAPI Referenceもしくは公式チュートリアルをご参照ください。

from datasets import load_dataset

# 例1：
# Dataset型としてロードする。

imagenet = load_dataset("timm/imagenet-1k-wds", split="train")

# 例2：
# IterableDataset型としてロードする。

iterable_imagenet = load_dataset("timm/imagenet-1k-wds", split="train", streaming=True)

すでにコード内で記述されたJSONなどをDataset（もしくはIterableDataset）オブジェクトに変換する場合は、from_xxx()を使用します。

どのような関数があるかはDatasetのAPI Referenceをご参照ください。（「 from_ 」と検索をかけると見つけやすいと思います。）

# 例1：
# Dataset型としてJSONから変換する。

my_dataset = Dataset.from_dict({"col_1": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]})

# 例2：
# IterableDataset型としてGeneratorから変換する。

def my_generator(n):

    for i in range(n):

        yield {"col_1": i}

my_iterable_dataset = IterableDataset.from_generator(my_generator, gen_kwargs={"n": 10})

Dataset.map関数

Dataset（もしくはIterableDataset）でもpandasなどと同じようにmap関数を使用できます。

もしもmap関数がピンとこないかたは別途map関数に関するネット記事などをお読みください。参考

Dataset.map()を使用するとDatasetオブジェクト内の全てのデータに対して即座に処理を実行することができます。

my_dataset = my_dataset.map(process_fn)

ここでprocess_fnは処理プロセスを定義した関数ですが、Dataset.mapのAPI Referenceを確認すると戻り値としては**Dict[str, Any]**が期待されていることが分かります。

def process_fn(examples): # exampleはDataset内の各要素です。
    text = example("text") # 例えば、Datasetの1つの要素内にtextが存在すると仮定します。
    text_head = text[:10] # なんらかの処理（ここでは先頭10文字を取得）を施します。
    return {"text_head": text_head}

processed_dataset = my_dataset.map(process_fn)

続いて、IterableDatasetのmap関数についてです。

既に何度か説明しているためご理解いただけていると思いますが、IterableDatasetはDatasetと違って全量を把握できていません。

そのためDatasetとは違ってmap関数を適用してもデータセット全体に対して即座に処理を行うことはありません。

チュートリアルのEager data processing and lazy data processingの部分を見てみましょう。

ここには以下の様に書かれています。

On the other hand, due to the “lazy” nature of an IterableDataset, calling IterableDataset.map() does not apply your map function over the full dataset. Instead, your map function is applied on-the-fly.

Because of that, you can chain multiple processing steps and they will all run at once when you start iterating over the dataset

これが何を言っているかまとめると次のようになります。

• map()を使用した段階ではデータセット全体に反映されない。（Datasetとは違う）

• on-the-fly（データを1つずつ取り出すとき）に適用されるため複数のprocess_fnを適用できる。

これによってIterable DatasetではDatasetとは違い、複数のprocess_fnを適用する場合でも無駄な中間データを生成することなく柔軟にメモリ効率よく処理を適用できます。

my_iterable_dataset = my_iterable_dataset.map(process_fn_1)

my_iterable_dataset = my_iterable_dataset.filter(filter_fn)

my_iterable_dataset = my_iterable_dataset.map(process_fn_2)

# process_fn_1, filter_fn and process_fn_2 are applied on-the-fly when iterating over the dataset

for example in my_iterable_dataset:  

    print(example)

    break

Tips

最後にTipsとして書かせていただきますが速度に関する重要な内容を抜粋します。

speed-differencesに書かれている内容なのですが

Regular Dataset objects are based on Arrow which provides fast random access to the rows. Thanks to memory mapping and the fact that Arrow is an in-memory format, reading data from disk doesn’t do expensive system calls and >deserialization. It provides even faster data loading when iterating using a for loop by iterating on contiguous Arrow record batches.

However as soon as your Dataset has an indices mapping (via Dataset.shuffle() for example), the speed can become 10x slower. This is because there is an extra step to get the row index to read using the indices mapping, and most >importantly, you aren’t reading contiguous chunks of data anymore. To restore the speed, you’d need to rewrite the entire dataset on your disk again using Dataset.flatten_indices(), which removes the indices mapping. This may take a lot of >time depending on the size of your dataset though

my_dataset[0]  # fast

my_dataset = my_dataset.shuffle(seed=42)

my_dataset[0]  # up to 10x slower

my_dataset = my_dataset.flatten_indices()  # rewrite the shuffled dataset on disk as contiguous chunks of data

my_dataset[0]  # fast again

ここではDataset.shuffle()を使用すると、その後の処理で10倍遅くなる可能性があると述べられています。

これの原因はshuffleした時に作成されるインデックスマッピングにあります。

インデックスマッピングとはなんぞやと言うと、データセットの元々のインデックスと実際にアクセスする順序（シャッフルされた順序）を対応づけるものです。

これが存在することによって、もとのデータセットの構造は変更することなくシャッフルすることができるのですがインデックスマッピングへの参照が処理に対して追加されるため遅くなるということです。

スピードを元の速度に戻すためには flatten_indices() を使用すればいいのですが、データセットの構造（順序）そのものが変更されてしまうことに注意してください。

終わりに

いかがだったでしょうか？
LLMのFinetuningと言うと、初学者の方のほとんど（私を含め）が真っ先にTransformerライブラリやUnslothなどに飛びつきがちですが、Finetuningにおいて最も重要なのは訓練に使用するデータセットそのものであり、そのデータセットを効率よく扱うためのdatasetsライブラリの基本的概念を理解するのは非常に重要なことだと考えています。

本記事が読者の皆様にとって有益なものであると幸いです。
間違いなどございましたらコメント欄でご指摘お願いいたします。