一、背景介紹

LLM 的模型參數(shù)量很大，其模型效果也需要巨大的語料庫支撐，LLM 預訓練需要的 Token 數(shù)已經(jīng)從早期的 300B Token 逐漸增加到 1.4T，甚至進一步擴展到 3T 以上。本文中我們具體介紹 LLM 預訓練語料庫的來源，構(gòu)建語料庫的預處理過程以及 LLM 預訓練的 Dataset 存儲、混合、加載方式。

二、常見語料庫

雖然不同 LLM 的模型大小不同，預訓練的 Token 數(shù)也各不一樣，但是其原始的語料都大同小異，主要有幾種類型：CommonCrawl、Wikipedia、Books、Code、ArXiv、Reddit links 等。

2.1 CommonCrawl

CommonCrawl 是一個免費、開放的網(wǎng)絡爬蟲數(shù)據(jù)集，旨在提供大規(guī)模的網(wǎng)頁抓取數(shù)據(jù)，使研究人員、開發(fā)者和數(shù)據(jù)科學家能夠訪問和分析互聯(lián)網(wǎng)上的信息。該數(shù)據(jù)集由 Common Crawl Foundation 維護，該基金會是一個非營利性組織，致力于促進網(wǎng)絡信息的開放共享。

CommonCrawl 數(shù)據(jù)集非常大，并且在不斷地更新中，具體可參考 Common Crawl - Overview，其中最新的 CC-MAIN-2023-50 共包含 3.35B 個網(wǎng)頁，壓縮后的數(shù)據(jù)超過 130TB。具體如下圖所示：

由于 CommonCrawl 數(shù)據(jù)集過于龐大，并且包含很多噪聲，處理的成本很高，因此也有其他研究者提供了相應處理過的子集，比如 C4（Colossal Clean Crawled Corpus），可以參考 GitHub - google-research/text-to-text-transfer-transformer。

2.2 Wikipedia

Wikipedia 是一個由全球志愿者維護的在線百科全書項目。其包含多種語言，涉及的領域也非常廣，并且質(zhì)量非常高。比如如下圖所示，“Large language model” 頁面有 29 種語言，并且分了各個部分進行介紹：

2.3 Books

書籍是另一種高質(zhì)量的語料庫，與其他語料庫相比，其涵蓋的語言、領域更加廣泛，內(nèi)容也更加正式?？傮w來說，使用書籍作為語料庫預訓練 LLM 可以獲得如下好處：

• 廣泛的知識涵蓋：書籍包含很多領域，比如科學、歷史、文學以及技術等。書籍能使 LLM 接觸豐富多樣的知識，有助于提高其對各種主題的理解和表達能力。
• 豐富的語言表達：書籍中通常包含豐富而復雜的語言表達，包括各種風格、修辭和專業(yè)術語。通過學習書籍中的語言使用方式，LLM 可以更好地捕捉到語境、上下文和語法結(jié)構(gòu)，提高其生成自然語言的能力。
• 專業(yè)的領域知識：一些書籍涉及特定領域的深度知識，如科學、法律、醫(yī)學等。在 LLM 的訓練中使用這些書籍可以使模型更好地理解和生成與這些領域相關的內(nèi)容。
• 多樣性的文本結(jié)構(gòu)：書籍中的文本結(jié)構(gòu)多種多樣，包括章節(jié)、段落、腳注等。通過訓練模型處理這些不同層次和結(jié)構(gòu)的文本，有助于提高其對復雜文檔和長文本的理解和處理能力。
• 知識結(jié)構(gòu)和推理能力：書籍中的內(nèi)容通常有一定的邏輯和知識結(jié)構(gòu)，通過訓練模型學習這些結(jié)構(gòu)，可以提高其在理解和生成邏輯推理、連貫性論述方面的能力。
• 語言多樣性：書籍中使用的語言可能涵蓋多種方言、俚語和文學風格，這有助于訓練模型更好地理解和生成多樣化的語言表達。

2.4 Code

當前很多 LLM 預訓練語料中也會加入 Code，比如來自 Github、Gitlab 或者編程問答網(wǎng)站（比如 StackOverflow）的語料，因為其不僅對 LLM 理解編程語言，代碼注釋和生成代碼很有幫助，也有研究表明其對 LLM 的推理能力至關重要。

2.5 ArXiv

ArXiv（??https://arxiv.org/??） 是一個包含各個學科領域的預印本（Preprint）平臺，涵蓋數(shù)學、物理、計算機科學等多個學科，包含大量領域特定的術語和語言，比如數(shù)學符號、專業(yè)術語等。在預訓練語料中加入 ArXiv 中的論文可以使 LLM 接觸到廣泛的學術知識、提高對不同學科的理解能力。

2.6 Stack Exchange

Stack Exchange （??https://stackexchange.com/??）是一個高質(zhì)量問答網(wǎng)站，涵蓋從計算機科學到化學等多個領域。

三、數(shù)據(jù)預處理

3.1 概述

收集到預料數(shù)據(jù)往往不會直接用于預訓練，因為其通常比較大，并且包含很多冗余和噪聲，需要進一步的過濾清理，此外，有些數(shù)據(jù)（尤其抓取的網(wǎng)頁數(shù)據(jù)）中還可能包含一些敏感信息，比如個人的身份信息、家庭信息以及其他色情或者敏感信息，都需要進一步的處理。

如下圖 Fig. 7 （來自 [2303.18223] A Survey of Large Language Models）所示，常見的數(shù)據(jù)處理包含質(zhì)量過濾（Quality Filtering）、去重（De-deplication）、隱私擦除（Privacy Reduction）、Tokenization、數(shù)據(jù)混合等：

3.2 LLaMA-1

LLaMA 是 Meta 發(fā)布的 LLM，如下圖所示為 LLaMA-1 中預訓練語料的來源、混合比例等統(tǒng)計信息，經(jīng) Tokenizer （采用 BPE 算法）后生成 1.4T Token：

針對不同的數(shù)據(jù)采用了不同的預處理策略：

• English CommonCrawl：使用 CCNet Pipeline 處理了 2017 年到 2020 年的 5 個 CommonCrawl 數(shù)據(jù)。預處理包括：按行進行重復數(shù)據(jù)刪除、基于fastText 語言分類器刪除非英語數(shù)據(jù)、使用 ngram 語言模型刪除低質(zhì)數(shù)據(jù)。
• C4：雖然 C4 數(shù)據(jù)集來自 CommonCrawl，但是作者發(fā)現(xiàn)使用不同預處理的 CommonCrawl 數(shù)據(jù)集可以提升性能，因此依舊使用了 C4。同時也應用了重復數(shù)據(jù)刪除、語言識別相關操作。
• Github：來自 Google 的 BigQuery，只保留符合 Apache、BSD 和 MIT 許可的項目。此外，根據(jù)行的長度或字母和數(shù)字的比例采用啟發(fā)式方法過濾低質(zhì)量文件。此外，也采用正則表達式刪除標題等內(nèi)容。最后，還在文件級按匹配度進行去重。
• Wikipedia：來自 2022 年 6月-8月版本，共 20 種語言。去除了超鏈接、注解和其它格式化樣式。
• Books：包含兩部分數(shù)據(jù)，來自 Gutenberg 項目的書籍，以及來自 The Pile 的 Books3 部分。按照書籍粒度進行去重，刪除超過 90% 重復內(nèi)容的書籍。
• ArXiv：作者以 Latex 格式處理 ArXiv 論文，并刪除第一節(jié)之前的所有內(nèi)容以及參考文獻，同時刪除 .tex 文件中的注釋、用戶自己編寫的定義和宏，以提高論文的一致性。
• Stack Exchange：作者僅保留 Stack Exchange 中 28 個最大的主題，并從文本中刪除 HTML 標簽，按分數(shù)（從高到低）對答案進行排序。

3.3 RefinedWeb

RefinedWeb 是阿布扎比的 TII 基于 CommonCrawl 構(gòu)建的語料庫，其數(shù)據(jù)集經(jīng)過嚴格的過濾和去重，總共包含 5T Token，不過只開源了 600B Token 的子集。同時還發(fā)布了 Falcon-7B 和 Falcon-40B，分別基于 1.5 T Token 和 1 T Token 訓練而來，其 80% 以上數(shù)據(jù)來自 RefinedWeb。

如下圖所示為 RefinedWeb 的構(gòu)建過程，其包含幾種主要的預處理過程：

• URL filtering：目的是為了過濾掉欺詐、色情、暴力和賭博等網(wǎng)站。作者首先收集了一個 4.6M 的黑名單網(wǎng)站列表，然后構(gòu)建了一個 URL 打分機制，以避免黑名單中的誤報，比如流行的博客、醫(yī)療、法律等頁面。此外，后續(xù)的預訓練中也要與其他優(yōu)質(zhì)語料庫混合，因此也過濾了高質(zhì)量的語料來源，比如Wikipedia和arXiv等。
• Text extraction：提取頁面中的主要內(nèi)容，忽略目錄、頁眉、頁腳和廣告等。
• Language identification：使用 CCNet 中的 fastText 語言分類器，可以將文檔分類為 176 種語言，作者刪除了排名靠前的語言中得分低于 0.65的文檔。此階段清理了一半左右的文檔。
• Repetition removal：作者基于啟發(fā)式方法刪除了重復的行、段落等內(nèi)容。
• Document-wise filtering：文檔中包含了很多機器生成的垃圾郵件等內(nèi)容，其不適于語言建模，作者通過總長度、字符和單詞比率等因素刪除了異常頁面。這個階段需要針對每種語言調(diào)整。
• Line-wise corrections：許多文檔的行中存在無關的交錯，比如社交媒體的點贊、導航按鈕等。因此，作者設計了一個行校正器，對行進行修正，同時如果修改超過文檔的 5%，則刪除整個文檔。
• Fuzzy deduplication：作者在文檔級別采用了 MinHash 來刪除相似的文檔。
• Exact deduplication：作者在序列級別采用精確字符串匹配的方式進一步對文檔進行去重，可以刪除特定的內(nèi)容，比如特定的免責聲明或通知等。

?

3.4 Baichuan 2

Baichuan 2 是百川發(fā)布的 LLM，其構(gòu)建了中英語料數(shù)據(jù)集，預訓練語料同樣包含網(wǎng)頁、書籍、論文和代碼。大致的數(shù)據(jù)分布如下圖 Figure 1 所示（圖片來自 [2309.10305] Baichuan 2: Open Large-scale Language Models）：

對于數(shù)據(jù)處理，作者聚焦在數(shù)據(jù)頻率和質(zhì)量上，數(shù)據(jù)頻率主要依賴聚類和去重。關于數(shù)據(jù)去重和聚類，Baichuan 2 采用基于 LSH 特征和稠密 Embedding 特征的方案。根據(jù)聚類，可以對單個文檔、段落和句子進行重復數(shù)據(jù)刪除以及打分，然后這些分數(shù)也可以用于預訓練中的數(shù)據(jù)采樣。

其整個數(shù)據(jù)預處理過程及各階段清理的數(shù)據(jù)占比如下圖 Figure 2 所示，其中灰色部分為刪除的數(shù)據(jù)比例：

Baichuan 2 的 Tokenizer 同樣采用 BPE（來自 SentencePiece），Tokenizer 后包含 2.6T Token，如下圖 Table 2 所示為 Baichuan 2 與其他模型的詞表大小及壓縮率比較：

3.5 Qwen

Qwen 是阿里發(fā)布的 LLM，其預訓練數(shù)據(jù)集包含網(wǎng)頁、百科全書、書籍以及代碼等，其數(shù)據(jù)集同樣是多語言的，主要是英文和中文。

為了保證預訓練數(shù)據(jù)的質(zhì)量，作者同樣采用了全面的預處理：

• 針對網(wǎng)頁數(shù)據(jù)，先從 HTML 中提取文本并使用語言識別工具來確定語言。
• 為了增加數(shù)據(jù)多樣性，同樣進行了數(shù)據(jù)去重，包括歸一化后的精確匹配去重和使用 MinHash 和 LSH 算法的模糊重復去重。
• 為了過濾低質(zhì)數(shù)據(jù)，作者結(jié)合了基于規(guī)則和基于機器學習的方法，具體來說，使用多種模型對內(nèi)容進行打分，包括語言模型、文本質(zhì)量評分模型以及識別潛在攻擊性和不當內(nèi)容的模型。
• 為了進一步提高數(shù)據(jù)質(zhì)量，作者有選擇地對某些數(shù)據(jù)進行上采樣，以確保模型接受更多高質(zhì)量數(shù)據(jù)。
• 此外，有些研究表明，在預訓練中使用多任務指令數(shù)據(jù)可以增強其 zero-shot 和 few-shot 數(shù)據(jù)，因此作者額外添加了高質(zhì)量的指令數(shù)據(jù)。

其 Tokenizer 遵照 GPT-3.5 和 GPT-4，同樣采用 BPE 方法，其詞表大小為 152K，最終生成 3T Token。

3.6 Skywork

Skywork-13B 是昆侖萬維的天工團隊發(fā)布的 LLM，其首先構(gòu)建了一個巨大的、高質(zhì)量的數(shù)據(jù)集 SkyPile，超過 6T Token，并開源了一個 150B Token 的子集，其原始語料為網(wǎng)頁數(shù)據(jù)，地址為 Skywork/SkyPile-150B · Datasets at Hugging Face。SkyPile 的數(shù)據(jù)包含多個來源，絕大部分是公開可訪問的。

Skywork 的數(shù)據(jù)處理和其他模型類似，包含幾個部分：

• Structural Extraction：數(shù)據(jù)集主要來源是網(wǎng)頁，因此第一階段的目標是提取相關內(nèi)容，同時刪除導航欄、特定站點的聯(lián)系信息等無關文本元素，保留連續(xù)的中長文本段落。
• Distribution Filtering：LLM 預訓練往往需要包含廣泛的領域知識，之前的模型通常是為每個文檔或網(wǎng)頁分配類別標簽，從而手動決定語料庫的組成。而本文中，作者摒棄了以標簽為中心的方法，核心是對文本片段之間的語義親和性進行測試，從而識別并刪除重復率高的文本塊。
• Deduplication：本文中作者把 Deduplication 當作 Distribution Filtering 的一部分。
• Quality Filtering：作者同樣使用 CCNet Pipeline 來執(zhí)行過濾任務，以刪除劣質(zhì)內(nèi)容和排除中文、英文以外的頁面。作者訓練了一個二分類器，來預測給定網(wǎng)頁適合作為 Wikipedia 參考的可能性。這一階段的結(jié)果會被組織成不同的質(zhì)量類別，并且只保留高質(zhì)量的組。

其 Tokenizer 同樣采用 BPE，詞表分布如下圖 Table 2 所示：

Skywork-13B 模型的預訓練語料包含 3.2T Token，從 SkyPile 采樣而來，其預訓練分為兩個階段，第一階段使用 2T Token，分布如下圖 Table 1 所示，第二階段采樣剩下的 1.2T Token：

3.7 DeepSeek

在 CC_Cleaner：一種絲滑高效且易擴展的數(shù)據(jù)清洗流程 中也介紹了幻方 AI LLM 的詳細數(shù)據(jù)集預處理流程，大體過程類似，這里不再具體介紹。

3.8 總結(jié)

從上可以看出，預訓練語料預處理基本上涵蓋幾個步驟：

• 過濾：

按 URL 過濾：按網(wǎng)站剔除欺詐、成人等網(wǎng)站頁面，同樣也可以刪除與其他數(shù)據(jù)集可能重復的頁面，比如 Wikipedia、arXiv 等。

按語言過濾（fastText 語言分類）：比如只保留英文和中文。

按質(zhì)量過濾：比如使用語言模型判斷序列的困惑度，然后刪除低質(zhì)數(shù)據(jù)。

• 去重：
• 文檔級去重：采用 MinHash 等方式刪除高度重復的文檔。
• 句子級去重：通過精確字符串匹配，或者基于 Embedding 的語義相似性刪除重復的語句。
• 交叉去重：預訓練語料可能來自不同的數(shù)據(jù)集，也需要交叉比對去重，比如從 CommonCrawl 中刪除 Wikipedia、arXiv 等數(shù)據(jù)；此外，通常也需要刪除與評估集重疊的數(shù)據(jù)，避免數(shù)據(jù)泄露。
• 隱私和許可：
• 隱私去除：為了避免隱私泄露，需要移除或者脫敏隱私數(shù)據(jù)，比如個人身份信息，姓名、身份證、住址等。
• 許可規(guī)范：在采集預訓練語料時需要遵循相應的許可協(xié)議，避免出現(xiàn)侵權(quán)等問題。比如代碼數(shù)據(jù)中只保留符合 Apache、BSD 和 MIT 許可的項目。
• 分詞：預訓練語料在輸入模型之前都需要經(jīng)過 Tokenizer 分詞，大部分模型都采用 BPE 算法。
• 數(shù)據(jù)混合：預訓練語料通常包含不同的數(shù)據(jù)集，有些質(zhì)量比較高，比如書籍、Wikipedia 等，然而其數(shù)據(jù)量可能不多，此時在真正訓練的時候可以給予更高的采樣權(quán)重。

四、數(shù)據(jù)存儲和加載

目前很多 LLM 預訓練會采用 NVIDIA 的 Megatron-LM 項目或者是 Microsoft 基于此改進的 DeepSpeed-Megatron 項目，其預訓練數(shù)據(jù)集的存儲格式和加載方式是一致的，此處我們以此為例進行介紹。

4.1 原始 Dataset 結(jié)構(gòu)

實際的預訓練語料在訓練之前都會先經(jīng)過 Tokenizer 分詞，轉(zhuǎn)換為 Binary 數(shù)據(jù)集（還沒有 shuffle 和采樣）。分詞后的數(shù)據(jù)都以 Token ID 的方式存儲，數(shù)據(jù)的大小基本等價于 Token 數(shù)乘以每個 Token ID 的字節(jié)數(shù)。

• 如果詞表大小比較小，小于 65536，則可以用 Uint16 表示，存儲占磁盤大小基本等于 2*Token 數(shù)。
• 很多中文 LLM 需要包含中文詞表，詞表數(shù)往往超過這個限制，需要使用 Uint32，導致數(shù)據(jù)大小變?yōu)?4*Token 數(shù)。

同時，數(shù)據(jù)集通常包含不同來源，比如 CommonCrawl，Wikipedia，而且數(shù)據(jù)集有大有小，為了避免單個數(shù)據(jù)集過大，會將數(shù)據(jù)集切分為不同的 part，每個 part 都相當于一個新的子集，但是來自同一個數(shù)據(jù)集的不同 part 需要有相同的權(quán)重。此外，每個 part 都有 idx 和 bin 兩個文件。如下所示為一些子數(shù)據(jù)集的示例：

• en-CommonCrawl-part18.idx
• en-CommonCrawl-part18.bin
• en-Wikipedia-part0.idx
• en-Wikipedia-part0.bin

其中 idx 文件對應索引文件，bin 對應實際的 Token ID，如下圖所示：

• Index：包含 Head 和 Buffer 兩部分（實際是連續(xù)的）

Head：存儲 magic、version、dtype、len 和 doc_count

Buffer: 存儲 Bin 中 Document 的起始位置和大小

• Bin：存儲實際的 Document，比如根據(jù)points[m]和sizes[m]即可以從 Bin 中獲得Document m。

?

需要說明的是，每個 Document 都已經(jīng) Token 化，并且已經(jīng)添加過起始 Token <s> 和終止 Token </s>。

4.2 GPT Dataset 結(jié)構(gòu)

在 Dataset Blending 混合階段可以計算獲得每個 GPT Dataset 數(shù)據(jù)需要生成多少個 Sample，根據(jù) Sample 數(shù)目也可以進一步計算得到該 Dataset 需要過幾個 Epoch。

如下圖所示：

• _num_tokens：根據(jù) Dataset 中每個 Document 的 size 累積即可獲得當前 Dataset 總的 Tokens 數(shù)目。
• _num_epochs：根據(jù)需要采樣的 num_samples，序列長度（seq_length），每個 Epoch 總的 Tokens 數(shù)目（tokens_per_epoch）即可以計算獲得數(shù)據(jù)需要過幾個 Epoch。需要說明的是，倒數(shù)第二行的（total_tokens - 1）是因為 Sample 中的 Token 個數(shù)會比 seq_length 多 1 個，但是多的這一個又會作為下一個 Sample 的起始 Token，因此總共需要的 Tokens 數(shù)目為 num_samples * seq_length + 1。

?

此外，Dataset 還需要有一定的 Shuffle 操作，總的來說有兩個 Shuffle：

• Dataset 中的不同 Document 會進行 Shuffle，對應doc_idx。
• Dataset 中的 Sample 也會 Shuffle，對應shuffle_idx。

如下圖所示：

• shuffle_idx：長度為 num_sample，存儲的是 shuffle 過的 Sample 索引，圖中的示例為 16 個 Sample，也就是實際訓練的時候先讀  Sample 5，然后讀 Sample 1，以此類推。
• doc_idx：根據(jù)原始數(shù)據(jù)可以知道數(shù)據(jù)集中有多少個 Document，然后根據(jù)該數(shù)據(jù)集需要重復的 Epoch 數(shù)目即可以得到總共需要有多少個 Document。圖中的示例假設 Dataset 中總共有 0, 1, 2, 3, 4, 5 這 6 個 Document，對應的 Epoch 為 4，需要說明的是，最后一個 Epoch 需要特殊處理，后續(xù)再介紹。
• sample_idx：長度為 num_sample + 1，因為其中存儲的是 Sample 的起始位置和終止位置，也就是該 Sample 對應了 doc_idx 中的哪些 Document。

比如 shuffle_idx 中的綠色 4這個 Sample 由 sample_idx 中第 4 個位置和第 5 個位置確定。

sample_idx 的第 4 個位置的 idx 對應起始的 Document 索引的位置，也就是綠色箭頭指向的 doc_idx 中的 3，而 offset 為 2，則表明要從Document 3的第 2+1=3 個 Token 開始。

sample_idx 的第 5 個位置的 idx 對應終止的 Document 索引的位置，也就是綠色箭頭指向的 doc_idx 中的 0，而 offset 為 3，則表明要從Document 0的第 3 個 Token 終止，但是每個 Sample 中在結(jié)束位置需要有一個額外的 Token，因此實際是從第 4 個 Token 終止。

• sample：根據(jù)上述的起始 Document idx 和 offset 以及終止 Document idx 和 offset 即可以獲取到最終的 Sample，其中最后一個 * Token 也會被下一個 Sample 復用。因為 GPT 訓練的目的是讓每一個位置預測下一個位置的 Token，因此input和target正好交錯 1 個 Token，這也是為什么 Sample 中的 Token 數(shù)目要比 seq_length 多 1。

?

如下圖所示為 GPTDataset 中獲取一個 Sample 的示例：

• 首先，從 shuffle_dix 中獲取 shuffle 過的 sample index
• 然后，從 sample_idx 中獲取第一個和最后一個Document 的 idx 的位置和 offset
• 最后，根據(jù) idx 的起始和終止位置可以獲得當前 Sample 由哪些 Document 組成，當然，其中的第一個和最后一個 Document 還需要根據(jù) offset 截斷。需要說明的是，有些 Document 比較長，因此有可能存在一個 Sample 來自同一個 Document，也就是doc_index_f 等于 doc_iindex_l。

?

如下所示，如果第一步不取 shuffle_idx，也就是直接打印真實的第 0,1,2 個 Sample，可以看出 Sample 的長度為 4097（seq_length==4096），并且每個 Sample 中的最后一個和下一個 Sample 的第一個重疊，比如 1670 和 10870。

如下圖所示為構(gòu)建 doc_idx 的過程，可以看出其最后一個 Epoch 進行了特殊處理，這是因為根據(jù) num_sample 算出來的 Epoch 很可能不是整數(shù)，也就是必然有些 Document 重復多，有些重復少，但需要盡可能地保證所有 Document 的采樣概率相同。

比如，Document [0, 1, 2, 3] 重復 2.5 次：

• 如果直接全局 shuffle 再截斷，比如從 [2, 0, 3, 2, 1, 3, 0, 3, 2, 0, 1, 1]，截斷后為 [2, 0, 3, 2, 1, 3, 0, 3, 2, 0]，這樣 Document 0,2,3 都出現(xiàn)了 3 次，而 Document 1 只出現(xiàn)了 1 次。
• 如果最后一個 Epoch 獨立 Shuffle，比如 [2, 0, 3, 3, 1, 0 ,2, 1,0,3,2,1]，此時截斷后為 [2, 0, 3, 3, 1, 0 ,2, 1,0,3]，可以保證采樣次數(shù)最多差 1。
• 目前還不確定為什么不每個 Epoch 都獨立 shuffle 拼接后再截斷。

?

實際上獲得的 doc_idx 并沒有經(jīng)過截斷，還是完整的 Epoch，因此在 shuffle_idx（也就是 shuffle Samples）時也需要特殊處理最后一個 Epoch：

4.3 GPT Dataset Blending

LLM 在訓練數(shù)據(jù)集中往往包含多個子數(shù)據(jù)集，實際訓練的時候每個數(shù)據(jù)集會使用不同的權(quán)重。比如說有如下三個數(shù)據(jù)集（數(shù)據(jù)集在之前的步驟中已經(jīng) shuffle）：

• A：100 Samples，采樣權(quán)重 0.3
• B：50 Samples，采樣權(quán)重 0.2
• C：400 Samples，采樣權(quán)重 0.5

假設需要訓練 1000 個 Samples，則相當于：

• A：使用 300 個 Samples，也就是數(shù)據(jù)過 3 輪
• B：使用 200 個 Samples，也就是數(shù)據(jù)過 4 輪
• C：使用 500 個 Samples，也就是數(shù)據(jù)過 1.25 輪

構(gòu)建混合數(shù)據(jù)集 index 的目標是構(gòu)建訓練需要的 1000 個 Samples 的索引，需要知道：

• Sample 來自哪個 Dataset：dataset_index
• 對應 Dataset 的第幾個 Sample：dataset_sample_index

對應計算過程如下：

• 遍歷每個 Sample 位置

根據(jù)當前位置之前已經(jīng)采樣過的 Sample 計算對應 Dataset 的權(quán)重，比如 A 為 0.34（+0.04），B 為 0.18（-0.02），C 為 0.45（-0.05），表明已采樣的 Sample 中 Dataset C 占比與設定目標相差更多，因此此處從 Dataset C 中采樣

C 中樣本計數(shù) +1（也就是如果最后把 dataset_sample_index 中來自 Dataset C 的數(shù)據(jù)拿出來，則按順序應該是 [0, 1, 2, …, 498, 499]）

更新 dataset_index 和 dataset_sample_index

需要說明的是，上述的 Sample 可以理解為拼接好的滿足 max_seq 的序列，來自同一個 Dataset，但可能來自不同的句子。

具體代碼位于 ??https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/main/megatron/data/helpers.cpp#L36-L97??，如下所示：

此外，原始 Megatron-DeepSpeed 中的 dataset_index 存儲的是 uint8 類型，也就是最多只能有 256 個 Dataset，但是實際上當前 LLM 預訓練的 Dataset 可能很多，比如有 1000 多個，此時有兩種方案：

1. 修改代碼，dataset_index 存儲改為 uint16 或 uint32。
2. 將 Dataset 合并到 256 個以內(nèi)，但是需要保證合并的 Dataset 的 Weight 相同，并且在 shuffle 前合并，不然不等價。

4.4 數(shù)據(jù)加載

如下圖所示，BlendableDataset 為實際 LLM 預訓練使用的 Dataset，其在初始化階段完成索引構(gòu)建（可以 Cache），訓練中直接遍歷相應的 Sample 即可（返回數(shù)據(jù)包含子數(shù)據(jù)集索引及在子數(shù)據(jù)集中的位置索引）：

五、參考鏈接

1. ??https://commoncrawl.org/overview??
2. ??https://data.commoncrawl.org/crawl-data/CC-MAIN-2023-50/index.html??
3. ??https://arxiv.org/abs/2303.18223??
4. ??https://www.high-flyer.cn/en/blog/cc_cleaner/??
5. ??https://arxiv.org/abs/2309.10305??
6. ??https://huggingface.co/datasets/Skywork/SkyPile-150B??
7. ??http://arxiv.org/abs/2310.19341??
8. ??https://github.com/NVIDIA/Megatron-LM??
9. ??https://github.com/microsoft/Megatron-DeepSpeed??
10. ??https://lifearchitect.ai/whats-in-my-ai/??

本文轉(zhuǎn)載自??AI閑談??，作者： AI閑談 ????