從頭預訓練一個 LLaMA 3 超級 mini 杯

作者：Xode

原文：https://zhuanlan.zhihu.com/p/695130168

整理：青稞 AI

這次打算用 Hugging Face 的 API 來寫一份預訓練大（小）模型的代碼，也就是用 Trainer 來做預訓練。由於只是想練習一下，因此打算選一個極小模型 + 小數據集。爲了貼近主流，於是打算預訓練一個 LLaMA 3——不過是超迷你版本，大小僅不到 20M。

想起來曾經看到過的微軟的工作 TinyStories，探索的是語言模型在多小的情況下還能流利地講故事，工作非常直白、有趣，剛好也契合我的練習想法，於是這次來複現一下。

代碼放在這裏了：GitHub - Mxoder/TinyStories: 從頭預訓練一隻超迷你 LLaMA 3——復現 TinyStories。

https://github.com/Mxoder/TinyStories

1. 前期準備

讓我們先來想一想大概需要做什麼。

首先是模型架構的選擇。原工作用的是 GPT Neo 架構（可以看他們的 config），這個算是很老的模型了，最初是 EleutherAI 用來複現追蹤 GPT-3 的工作的，現在用的也比較少了。我打算選用 LLaMA 架構，也算是符合研究主流、便於推廣。LLaMA 3 主要多了個 GQA，也是現在模型的主流，我這裏也用一下。

其次是數據的選擇。既然是復現，就直接貫徹拿來主義，用原工作開源的數據集（主要是從頭生成要花不少 api 費用）。原工作第一版的時候用的是 GPT-3.5 生成的數據，後面社區有人更新了第二版，是用 GPT-4 生成的，比原數據更好，就用它了。

最後是訓練。其實我手上就兩張 3060 12G 和 4060 Ti 16G，訓這個確實是綽綽有餘，但我還是不想在桌前吵我自己，於是繼續用 Colab。現在 Colab 可以直接看到剩餘使用時長了，雖然已經被砍到只有 3h 左右的用卡時間，但至少心裏有個底，況且 3h 訓我們這個也完全夠了。

我們這次用到的 Hugging Face 的庫如下：

transformersacceleratedatasets

理論上比較新的版本都沒問題，但如果你很久沒更新了，最好用 pip install -U 來升級一下。

我這裏的用到的庫版本如下，供參考：

torch==2.2.1transformers==4.40.0accelerate==0.29.3datasets==2.18.0

另外，接下來的步驟講解主要是以 jupyter notebook 的形式展開的，並不是 .py 文件的形式，也就是說前面執行的變量會在中間儲存下來。

2. 原工作簡介

雖然是練習，但既然打着復現工作的名頭，還是來簡要回顧一下原工作究竟做了什麼吧。

原工作探索的問題是語言模型（LM）在文本連貫性上的表現。像早期的一些語言模型如 GPT-2，即使在一些 Common Craw 這樣的語料庫上大量預訓練後，也很難生成長的、連貫的文本。比如前幾年有一種 AI 玩具類型是做文本續寫，例如彩雲小夢，可以寫寫作文、小說什麼的，如果大家玩過就知道效果其實一言難盡，和今天的大模型完全沒法比，其實這就是 GPT-2 level 的續寫能力。

作者就在想，會不會是因爲訓練的語料庫太多、太寬泛，需要學習各種語法元素、詞彙、知識、推理等等，才導致小語言模型（SLM）沒法有一個很好的表現。作者決定專注於一個任務——短篇故事續寫，來探索一下 LM 的性能邊界。

作者用 GPT-4 和 GPT-3.5 構建了一個英文短篇小說數據集 TinyStories，將內容限制在三四歲兒童也能輕鬆理解的程度，並且使用不同的關鍵詞來讓故事的主題足夠豐富。此外，他們還加入了額外的關鍵詞，來控制故事有更曲折的走向、不同的結局等等。

作者用的模型基座架構是 GPT Neo，詞表大小約爲 50k，並且他們嘗試了不同的模型參數，調整了隱藏層維度（hidden_size）、隱藏層數（num_hidden_layers）等，來探索不同參數對於模型性能的影響。

作者的評估方式是經典的 GPT-4 監督打分模式，就是讓不同的 SLM 根據提示生成故事，然後 GPT-4 從設定好的不同維度來打分，主要有 Creativity、Grammar、Consistency 三項，分別代表創造性、語法正確性、上下文一致性。此外，作者額外加入了一套 TinyStories-Instruct 數據集，來訓練一批指令微調的 SLM，並測試他們的指令跟隨能力，也就是第四項 Instruct。

作者主要和 GPT-Neo 以及 GPT-2 的小中大杯進行了對比。

3. 模型初始化

讓我們正式開始復現！

3.1 決定模型的參數

首先是定義我們自己的模型。由於 LLaMA 3 的架構早就集成於 transformers 庫中，因此我們可以直接用 AutoConfig 初始化一個模型配置，傳入參數 model_type='llama' 即可。

架構確定了，那麼現在來探討一下模型具體參數，比如隱藏層大小、隱藏層數等等。我們先來看看 TinyStories 原工作的實驗結果：

可以看到，隱藏層維度從 64 增長到 256 時的收益是比較大的，往後收益就逐漸放緩了。而層數的影響並不如隱藏層維度那麼大，大而淺的網絡也能有不錯的表現（例如 hidden_size=1024, num_hidden_layers=1 的模型）。綜合考慮，我這裏選擇 hidden_size=256 和 num_hidden_layers=4。

其他參數方面，我們遵循現在主流的研究表現，將 FFN 的維度從傳統的 4 倍隱藏層維度設爲 8/3 倍（按 128 向上取整）。頭的數目我們設爲 16，並應用 GQA 機制。GQA 的實現在 transformers 中非常簡單，只需要配置 num_key_value_heads 即可。num_key_value_heads 取值和 num_attention_heads 相同時即爲 MHA 機制，取值爲 1 時即爲 MQA 機制。

綜上，我們的配置如下：

# 模型配置from transformers import AutoConfighidden_size = 256# 中間層取 8/3 倍，按 128 向上取整intermediate_size = (int(hidden_size * 8/3 / 128) + 1) * 128# 只改動我們需要調整的參數，其餘保持不變config = AutoConfig.for_model(    model_type='llama',    hidden_size=hidden_size,    intermediate_size=intermediate_size,    num_attention_heads=16,    num_hidden_layers=4,    num_key_value_heads=8                  # 分爲 8 組)'''LlamaConfig {  'attention_bias': false,                 # 不使用注意力偏置  'attention_dropout': 0.0,                # 注意力層的 dropout 比例  'bos_token_id': 1,                       # bos_token (begin of sentence) 的 id  'eos_token_id': 2,                       # eos_token (end of sentence) 的 id  'hidden_act': 'silu',                    # 隱藏層激活函數類型，silu 即 SwiGLU  'hidden_size': 256,                      # 隱藏層維度大小  'initializer_range': 0.02,               # 權重初始化範圍，會被後面的 Kaiming 初始化覆蓋  'intermediate_size': 768,                # 中間層大小，採用 8/3 倍而非 4 倍  'max_position_embeddings': 2048,  'model_type': 'llama',  'num_attention_heads': 16,  'num_hidden_layers': 4,  'num_key_value_heads': 8,  'pretraining_tp': 1,  'rms_norm_eps': 1e-06,  'rope_scaling': null,  'rope_theta': 10000.0,  'tie_word_embeddings': false,            # 頭尾 embedding 和 lm_head 是否共享權重  'transformers_version': '4.40.0',  'use_cache': true,  'vocab_size': 32000}'''

3.2 分詞器 Tokenizer

我這裏選用 LLaMA 2 的分詞器，因爲二代的詞表比較小（32k），LLaMA 3 的詞表太大了（128k），在 SLM 中會佔用太多的參數比重，並且這只是個專有任務數據訓練，沒必要用太大的詞表。

# 分詞器from transformers import AutoTokenizertokenizer = AutoTokenizer.from_pretrained('NousResearch/Llama-2-7b-hf')'''LlamaTokenizerFast(name_or_path='NousResearch/Llama-2-7b-hf', vocab_size=32000, model_max_length=1000000000000000019884624838656, is_fast=True, padding_side='left', truncation_side='right', special_tokens={'bos_token': '<s>', 'eos_token': '</s>', 'unk_token': '<unk>', 'pad_token': '<unk>'}, clean_up_tokenization_spaces=False),  added_tokens_decoder={    0: AddedToken('<unk>', rstrip=False, lstrip=False, single_word=False, normalized=True, special=True),    1: AddedToken('<s>', rstrip=False, lstrip=False, single_word=False, normalized=True, special=True),    2: AddedToken('</s>', rstrip=False, lstrip=False, single_word=False, normalized=True, special=True),}'''

另外注意這裏 padding_side='left'，如果不是的話需要設置 tokenizer.padding_side='left'，即批量填充的時候從左邊開始填充，這對於 decoder-only 的模型做生成任務是必要的，因爲我們本質上做的是 next token prediction，如果 pad 擋在了生成序列的右邊，會影響到模型生成。

# 假設 pad_token 就是 eos_token(</s>)# 從右邊填充Once upon a time </s></s></s></s>...# 從左邊填充</s></s></s></s>Once upon a time ...

3.3 模型實例化

接下來就是實例化模型，這裏就不用從預訓練模型加載 from_pretrained() 了，而是從配置加載 from_config()：

# 模型import torchfrom transformers import AutoModelForCausalLM# 能用 cuda 就用 cudadevice = 'cuda' if torch.cuda.is_available() else 'cpu'# 從配置加載模型model = AutoModelForCausalLM.from_config(                        config,    torch_dtype=torch.float32   # 全精度訓練).to(device)                    # 遷移到 device 上'''LlamaForCausalLM(  (model): LlamaModel(    (embed_tokens): Embedding(32000, 256)    (layers): ModuleList(      (0-3): 4 x LlamaDecoderLayer(        (self_attn): LlamaSdpaAttention(          (q_proj): Linear(in_features=256, out_features=256, bias=False)          (k_proj): Linear(in_features=256, out_features=128, bias=False)          (v_proj): Linear(in_features=256, out_features=128, bias=False)          (o_proj): Linear(in_features=256, out_features=256, bias=False)          (rotary_emb): LlamaRotaryEmbedding()        )        (mlp): LlamaMLP(          (gate_proj): Linear(in_features=256, out_features=768, bias=False)          (up_proj): Linear(in_features=256, out_features=768, bias=False)          (down_proj): Linear(in_features=768, out_features=256, bias=False)          (act_fn): SiLU()        )        (input_layernorm): LlamaRMSNorm()        (post_attention_layernorm): LlamaRMSNorm()      )    )    (norm): LlamaRMSNorm()  )  (lm_head): Linear(in_features=256, out_features=32000, bias=False))'''

可以看到，k_proj 和 v_proj 的 out_features 從 256 變爲了 128，這即是 GQA 機制。

此時，模型已經初始化了，讓我們來打印一下看看參數：

# 打印模型的每一層及其參數大小def print_model_parameters(model):    print('Layer Name & Parameters')    print('----------------------------')    total_params = 0    for name, parameter in model.named_parameters():        param_size = parameter.size()        param_count = torch.prod(torch.tensor(param_size)).item()        total_params += param_count        print(f'{name:50} | Size: {str(param_size):30} | Count: {str(param_count):20}')    print('----------------------------')    print(f'Total Parameters: {total_params} ({total_params / 1000000:.1f} M)')print_model_parameters(model)

得到結果如下：

Layer Name & Parameters----------------------------model.embed_tokens.weight                          | Size: torch.Size([32000, 256])       | Count: 8192000     model.layers.0.self_attn.q_proj.weight             | Size: torch.Size([256, 256])         | Count: 65536       model.layers.0.self_attn.k_proj.weight             | Size: torch.Size([128, 256])         | Count: 32768       model.layers.0.self_attn.v_proj.weight             | Size: torch.Size([128, 256])         | Count: 32768       model.layers.0.self_attn.o_proj.weight             | Size: torch.Size([256, 256])         | Count: 65536       model.layers.0.mlp.gate_proj.weight                | Size: torch.Size([768, 256])         | Count: 196608     model.layers.0.mlp.up_proj.weight                  | Size: torch.Size([768, 256])         | Count: 196608     model.layers.0.mlp.down_proj.weight                | Size: torch.Size([256, 768])         | Count: 196608     中間省略...model.layers.3.input_layernorm.weight              | Size: torch.Size([256])              | Count: 256         model.layers.3.post_attention_layernorm.weight     | Size: torch.Size([256])              | Count: 256         model.norm.weight                                  | Size: torch.Size([256])              | Count: 256         lm_head.weight                                     | Size: torch.Size([32000, 256])       | Count: 8192000     ----------------------------Total Parameters: 19532032 (19.5 M)

可以看到，我們的模型只有不到 20M！非常非常小，並且其中 Embedding 佔了大頭。

儘管模型還沒有訓練，但我們仍然可以測試一下推理：

def inference(    model: AutoModelForCausalLM,    tokenizer: AutoTokenizer,    input_text: str = 'Once upon a time, ',    max_new_tokens: int = 16):    inputs = tokenizer(input_text, return_tensors='pt').to(device)    outputs = model.generate(        **inputs,        pad_token_id=tokenizer.eos_token_id,        max_new_tokens=max_new_tokens,        do_sample=True,        top_k=40,        top_p=0.95,        temperature=0.8    )    generated_text = tokenizer.decode(        outputs[0],        skip_special_tokens=True    )    # print(outputs)    print(generated_text)inference(model, tokenizer)'''Once upon a time, Hostย crimeine /\ könnenlinewidth measurementresol perfectly Taylor measèresiones assetviron'''

嗯，的確是胡言亂語呢，不過可以正常推理，說明模型沒問題！

但現在模型是隨機初始化的，爲了讓模型更好地收斂，我們最好給模型一個更好的初始化方法，我這裏選用 Kaiming 初始化，比較適用於 ReLU 類的激活，當然也可以選用高斯初始化、Xavier 初始化等等。

# Kaiming 初始化def kaiming_initialization(model):    for name, param in model.named_parameters():        if 'weight' in name and param.dim() > 1:            torch.nn.init.kaiming_uniform_(param, mode='fan_in', nonlinearity='leaky_relu')        elif 'bias' in name:            # 一般偏置項可以初始化爲 0            torch.nn.init.constant_(param, 0)kaiming_initialization(model)

現在，我們的模型真正初始化完成了！如果你願意，可以先將這個初始化好的模型保存到本地，用 save_pretrained() 即可。

4. 數據集

讓我們繼續！

4.1 加載數據集

我們接下來需要從 Hugging Face 加載數據集，我這裏是建立在網絡暢通的基礎上的，如果你沒有用 Colab 或者網絡無法直連 Hugging Face，那麼也可以先下載到本地某個文件夾中，load_dataset 也可以直接讀取本地文件夾。我們要用的數據集路徑如下：noanabeshima/TinyStoriesV2 · Datasets at Hugging Face

https://huggingface.co/datasets/noanabeshima/TinyStoriesV2
# 加載數據集from datasets import load_datasetdataset_name_or_path = 'noanabeshima/TinyStoriesV2'        # 可以替換爲本地文件夾路徑# ds_train = load_dataset(dataset_name_or_path, split='train')        # 取全部數據ds_train = load_dataset(dataset_name_or_path, split='train[:10%]')    # 只取前 10 %，約 270k 條ds_val = load_dataset(dataset_name_or_path, split='validation')print(ds_train)print(ds_val)'''Dataset({    features: ['text'],    num_rows: 271769})Dataset({    features: ['text'],    num_rows: 27629})'''

我們來看看數據長什麼樣子：

# 查看前兩條print(ds_train[:2])'''{'text': ['Once upon a time, there was a reliable otter named Ollie. He lived in a river with his family. They all loved to play and swim together.\nOne day, Ollie\'s mom said, 'Ollie, hurry and get some fish for dinner!' Ollie swam fast to catch fish. He saw his friend, the duck. 'Hi, Ollie!' said the duck. 'Hi, duck!' said Ollie. 'I need to hurry and catch fish for my family.'\nWhile Ollie was catching fish, he found a big shiny stone. He thought, 'This is not a fish, but it is so pretty!' Ollie took the shiny stone home to show his family. They all looked at the shiny stone and smiled. The shiny stone made everyone happy, and they forgot about the fish for dinner.',  'One day, a little boy named Tim went to the park. He saw a big tiger. The tiger was not mean, but very easy to play with. Tim and the tiger played all day. They had lots of fun.\nThen, something unexpected happened. The tiger started to shake. Tim was scared. He did not know what was going on. But then, the tiger turned into a nice dog. Tim was very surprised.\nTim and the dog played together now. They were very happy. The dog was easy to play with too. At the end of the day, Tim went home with his new friend.']}'''

這裏需要注意，datasets 加載後的數據是 Dict[str, List[str]] 的形式的，並非 List[Dict[str, str]]。

# 數據長這樣{    'text': [        <text_1>, <text_2>, <text_3>, ...    ]}# 不是下面這樣[    {'text': <text_1>},    {'text': <text_2>},    {'text': <text_3>},    ...]

4.2 數據預處理

接下來，我們要將數據預處理一下，也就是用 tokenizer 進行 tokenize。讓我們來寫一個處理函數：

from typing import Dict, Listdef process_func(    examples: Dict[str, List]) -> Dict[str, List]:    max_token = 2048    # 設置最長 token 數目，對於我們當前任務，2048 絕對不會超    encoded_texts = tokenizer(examples['text'], add_special_tokens=False)    input_ids_list = encoded_texts['input_ids']    new_input_ids_list, new_attn_mask_list = [], []    for input_ids in input_ids_list:        temp = input_ids[-max_token+1:] + [tokenizer.eos_token_id]        new_input_ids_list.append(temp)        new_attn_mask_list.append([1] * len(temp))    return {        'input_ids': new_input_ids_list,        'attention_mask': new_attn_mask_list    }

我們來解析一下其中的一些點：

tokenizer 的 encode

encoded_texts = tokenizer(examples['text'], add_special_tokens=False)

根據前面的示例，我們知道這裏 examples['text'] 其實是一個 List[str]，當一個 List 傳入 tokenizer() 時，tokenizer 會自動進行 batch encode，得到的是 {'input_ids': List[int], 'attention_mask': List[int]}（當然，如果設置了 return_tensors='pt' 就會得到 Tensor）。

add_special_tokens=False 則是讓 tokenizer 不要加上特殊 token，在 LLaMA 中就是不會在句首加上 bos_token 。

text = 'Hello, world!'tokenizer(text)# {'input_ids': [1, 15043, 29892, 3186, 29991], 'attention_mask': [1, 1, 1, 1, 1]}tokenizer(text, add_special_tokens=False)# {'input_ids': [15043, 29892, 3186, 29991], 'attention_mask': [1, 1, 1, 1]}# 上面多了一個 1，即 tokenizer.bos_token_id，在 LLaMA 中對應的就是 <s>

填充還是截斷？

temp = input_ids[-max_token+1:] + [tokenizer.eos_token_id]new_input_ids_list.append(temp)new_attn_mask_list.append([1] * len(temp))

在這裏，我採用直接截斷的方式，最大截取當前輸入序列的後 (max_token - 1) 位，再加上一個 eos_token_id，組成總長度不超過 max_token 的序列。attention_mask 的長度保持一致，全爲 1。

這裏利用到了 list 的切片特性，input_ids[-max_token+1:] 可以獲取 min(max-token, len(input_ids)) - 1 的序列。

當然，也可以採取將超出長度部分再按照 max_token 來分塊，重新組裝。

應用在所有數據上

接下來，我們用 map() 函數，來將 process_func() 應用到 ds_train 和 ds_val 中的每個樣本：

num_proc = 8                                    # 處理數據時所用的線程數ds_train = ds_train.shuffle()                    # 訓練集打亂一下ds_train = ds_train.map(    process_func,    batched=True,    num_proc=num_proc,    remove_columns=ds_train.column_names,    desc='Running tokenizer on train_set: ')ds_val = ds_val.map(    process_func,    batched=True,    num_proc=num_proc,    remove_columns=ds_val.column_names,    desc='Running tokenizer on val_set: ')print(ds_train)print(ds_val)'''Dataset({    features: ['input_ids', 'attention_mask'],    num_rows: 271769})Dataset({    features: ['input_ids', 'attention_mask'],    num_rows: 27629})'''

數據預處理成功！

4.3 數據批處理——DataCollator

4.3.1 兩行代碼

我們在訓練的時候往往不會一條一條訓練，而是成批次地訓練，那麼我們就需要對數據做批處理。因此我們需要用到 transformers 中的一個工具系列——DataCollator[1]。

既然是預訓練，那麼就是讓模型在語料上做無監督學習，也就是我們熟知的 next token prediction，即根據前面的所有輸入來預測下一個 token，然後把新的 token 拼接在已有輸入上作爲下一輸入，如此往復，直到觸發停止設定（例如觸發 max_new_tokens）。

所以我們的訓練目標——或者說是 label——顯而易見，就是把輸入偏移一位當作預測目標，我們計算的就是輸出和這個目標之間的 loss：

...   once  upon  a     time   there  was  ...        # inputonce  upon  a     time  there  was    ...  ...        # label# label 錯開一位，是 input 的下一預測目標，計算的就是 input 和 label 之間的 loss

所以我們只需要把 input_ids 複製一份、再偏移一位，就可以作爲 labels 了。

…… 再等等，讓我們看看這條問題：Shifting ids to the right when training GPT-2 on text generation? - Beginners - Hugging Face Forums[2]：

這裏 sgugger 提到，在 Hugging Face 的實現裏，training 時已經實現好了偏移一位的邏輯，不需要我們再手動實現了。我們也可以在 transformers 的源碼裏看到這一點，例如 LLaMA[3] 的實現。

所以，我們只需要將 input_ids 直接複製一份作爲 labels 即可。

那麼怎麼做呢？我們可以用 DataCollatorForLanguageModeling，並設置 mlm=False：

from transformers import DataCollatorForLanguageModelingdata_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

兩行代碼，非常簡單！

不過我們可以稍微多講一點，這個 data collator 是如何發揮作用的？爲什麼選的是它而不是在微調中更常見的 DataCollatorForSeq2Seq？

4.3.2 More things……

實際上，就像 mlm 這個參數所顯示的一樣，DataCollatorForLanguageModeling 一開始其實是設計給 Bert 的 MLM 任務的。MLM 任務就是 Masked Language Modeling，掩碼語言建模，就是在一整個序列中挑選一部分 token 用 [mask] 給蓋住，讓模型去根據上下文預測被蓋住的是什麼 token。

# MLM 任務示意今 天 早 上 下 [mask] 了     # input今 天 早 上 下   雨   了     # label

可以看到，這樣建立的 label 就是原來的 input 的 copy，它是將 input 中隨機 mask 一部分，label 不變。

這幾乎就是我們想要的——只是不需要 mask。所以，我們設置 mlm=False 後，就可以直接得到 input_ids 的 copy 了。

讓我們繼續看看 DataCollatorForLanguageModeling 怎麼作用的：

# DataCollatorForLanguageModeling# 這裏的 tokenizer 選用的是 Qwen1.5 的，並非 LLaMA 的，只是做一個示意dc = DataCollatorForLanguageModeling(tokenizer, mlm=False)data = ['南京', '南京市', '南京市長江']raw_tokens = [tokenizer(text) for text in data]print(f'tokenizer.pad_token_id: {tokenizer.pad_token_id}\n')print(dc(raw_tokens))'''tokenizer.pad_token_id: 151643{    'input_ids': tensor([[151643, 151643, 102034],                         [151643, 151643, 112891],                         [102034, 102975,  69177]]),    'attention_mask': tensor([[0, 0, 1],                              [0, 0, 1],                              [1, 1, 1]]),    'labels': tensor([[  -100,   -100, 102034],                      [  -100,   -100, 112891],                      [102034, 102975,  69177]])}'''

可以看到：

那麼微調裏常用的 DataCollatorForSeq2Seq 又是如何作用的呢？我們仍然用剛剛的數據例子：

# DataCollatorForSeq2Seq# 這裏的 tokenizer 選用的是 Qwen1.5 的，並非 LLaMA 的，只是做一個示意dc = DataCollatorForSeq2Seq(tokenizer)data = ['南京', '南京市', '南京市長江']raw_tokens = [tokenizer(text) for text in data]print(f'tokenizer.pad_token_id: {tokenizer.pad_token_id}\n')print(dc(raw_tokens))'''tokenizer.pad_token_id: 151643{    'input_ids': tensor([[151643, 151643, 102034],                         [151643, 151643, 112891],                         [102034, 102975,  69177]]),    'attention_mask': tensor([[0, 0, 1],                              [0, 0, 1],                              [1, 1, 1]])    # 注意沒有 labels 字段}'''

可以發現，DataCollatorForSeq2Seq 和 DataCollatorForLanguageModeling 一樣，做了批處理和 padding，但是沒有標籤 labels。原因是：DataCollatorForSeq2Seq 設計之初用於的任務和它的名字一樣，是序列到序列（seq2seq）任務，放到文本任務上，就是要有兩個 seq：輸入 text 和 輸出 label，比如下面的例子：

# DataCollatorForSeq2Seq# 這裏的 tokenizer 選用的是 Qwen1.5 的，並非 LLaMA 的，只是做一個示意dc = DataCollatorForSeq2Seq(tokenizer, padding=True)data = [('南京', '市長江大橋'), ('南京市', '長江大橋'), ('南京市長江', '大橋')]features = []for text, label in data:    feature = tokenizer(text)    feature['labels'] = tokenizer(label)['input_ids']    features.append(feature)print(f'tokenizer.pad_token_id: {tokenizer.pad_token_id}\n')print(dc(features))'''tokenizer.pad_token_id: 151643{    'input_ids': tensor([[151643, 151643, 102034],                         [151643, 151643, 112891],                         [102034, 102975,  69177]]),    'attention_mask': tensor([[0, 0, 1],                              [0, 0, 1],                              [1, 1, 1]]),    'labels': tensor([[102975,  69177, 106936],                      [  -100, 104924, 106936],                      [  -100,   -100, 106936]])}'''

可以看到：

5. 超迷你 LLaMA，啓動！

5.1 配置訓練參數

我們需要用到 transformers 的 TrainingArguments 來配置訓練參數，具體參數說明可以看這裏。

from transformers import TrainingArgumentstraining_args = TrainingArguments(    output_dir='saves',                         # 輸出路徑，包括模型檢查點、中間文件等    overwrite_output_dir=True,                  # 是否覆寫 output_dir    do_train=True,                              # 是否做訓練    do_eval=True,                               # 是否做評估    eval_steps=1000,                            # 評估步驟間隔    per_device_train_batch_size=4,              # 每設備批次    gradient_accumulation_steps=1,              # 梯度累計步大小，省顯存，但小模型沒必要，用 1 收斂比較快    learning_rate=1e-4,                         # 學習率大小    lr_scheduler_type='cosine',                 # 學習率調度策略，LLM 訓練一般都用餘弦    bf16=torch.cuda.is_bf16_supported(),        # 嘗試配置 bf16    fp16=not torch.cuda.is_bf16_supported(),    # bf16 不行就上 fp16    logging_steps=50,                           # 打印步驟間隔    report_to=None,                             # 日誌輸出目標，不想用 wandb 可以設置爲 None    num_train_epochs=2,                         # 訓練輪數，2 ~ 3 即可    save_steps=1000,                            # 檢查點保存步驟間隔    save_total_limit=2,                         # output_dir 內留存的檢查點最大數目    seed=3407                                   # 隨機種子)

如果你之前用了 wandb，現在想禁用掉，可以設置環境變量：

import osos.environ['WANDB_DISABLED'] = 'true'

5.2 配置 Trainer

同樣地，具體參數說明可以看這裏。

from transformers import Trainertrainer = Trainer(    model=model,                    # 模型實例    args=training_args,             # 訓練參數    train_dataset=ds_train,         # 訓練集    eval_dataset=ds_val,            # 驗證集（評估集）    tokenizer=tokenizer,            # 分詞器    data_collator=data_collator,    # data collator)

5.3 訓練與保存

配置好 Trainer 後，通過下列代碼即可啓動訓練：

trainer.train()

接下來只需要等待訓練完成。我用一個半小時訓練了 2 epochs，loss 達到了 1.6 左右。

訓練完成後，如果用的是 jupyter notebook，那麼此時 model 已經是訓練好的狀態了。我們可以再次推理試試看：

inference(    model,    tokenizer,    'Once upon a time, in a beautiful garden, there lived a little rabbit named Peter Rabbit.',    max_new_tokens=256)

得到如下結果：

Once upon a time, in a beautiful garden, there lived a little rabbit named Peter Rabbit. Peter had a friend named Rosie. They loved to play together. They would run, jump, and laugh all day long.One day, Robby saw a big box in his yard. He was curious and wanted to know what was inside. So, he went to his friend's house and asked, 'What are you doing, Spark?' May replied, 'I am making this big box in the garden, and I am trying to open it!'Timmy and Hopper went to find the big box. They found a key under a tree. They opened the box and found many toys inside. They were so happy to have a fun day with their new friend. They played with the toys all day long. And from that day on, whenever Ellie was a part of something, they would always remember the day they met by the big pond.

可以看到：

總之，這個超迷你 LLaMA 3 確實訓練完成了！我們可以將它保存到本地：

model_path = '...'model.save_pretrained(model_path)

也可以推送到 Hugging Face：

from huggingface_hub import notebook_loginrepo_name = 'TinyStories-LLaMA2-20M-256h-4l-GQA'notebook_login()    # 輸入 Access Tokensmodel.push_to_hub(repo_name)tokenizer.push_to_hub(repo_name)

6. 結尾

這次嘗試用 Trainer 來做一個模型的預訓練，以往都是用 Trainer 來做微調，這次也算是學習了一下吧。TinyStories 這個工作之前就有關注過，但一直沒顧上來複現一下，這次也算是簡單復現了個小模型出來，和原工作的豐富度確實是比不了，但也算完成一個 todo。

後面這個小模型可以繼續做 SFT，也就是做指令微調，可以和原工作一樣，給定故事背景、關鍵詞、開頭讓小模型續寫，也可以遷移到別的任務。不過由於我們的預訓練任務只針對了講短篇故事這一類任務，加上參數又特別少，如果直接遷移其它指令任務估計表現不會很好。

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/DJ121XtdffyAgyFxoAu7nQ

猜你喜歡