链载Ai

标题: 从头预训练一只超迷你 LLaMA 3 [打印本页]

作者: 链载Ai 时间: 13 小时前
标题: 从头预训练一只超迷你 LLaMA 3

ingFang SC', 'Hiragino Sans GB', 'Microsoft YaHei UI', 'Microsoft YaHei', Arial, sans-serif;margin: 1.5em 8px;letter-spacing: 0.1em;color: rgb(63, 63, 63);">这次打算用 Hugging Face 的 API 来写一份预训练大（小）模型的代码，也就是用 Trainer 来做预训练。由于只是想练习一下，因此打算选一个极小模型 + 小数据集。为了贴近主流，于是打算预训练一个 LLaMA 3——不过是超迷你版本，大小仅不到 20M。

ingFang SC', 'Hiragino Sans GB', 'Microsoft YaHei UI', 'Microsoft YaHei', Arial, sans-serif;margin: 1.5em 8px;letter-spacing: 0.1em;color: rgb(63, 63, 63);">想起来曾经看到过的微软的工作 TinyStories，探索的是语言模型在多小的情况下还能流利地讲故事，工作非常直白、有趣，刚好也契合我的练习想法，于是这次来复现一下。

ingFang SC', 'Hiragino Sans GB', 'Microsoft YaHei UI', 'Microsoft YaHei', Arial, sans-serif;margin: 1.5em 8px;letter-spacing: 0.1em;color: rgb(63, 63, 63);">代码放在这里了：GitHub - Mxoder/TinyStories: 从头预训练一只超迷你 LLaMA 3——复现 TinyStories。

ingFang SC', 'Hiragino Sans GB', 'Microsoft YaHei UI', 'Microsoft YaHei', Arial, sans-serif;overflow-x: auto;border-radius: 8px;margin: 10px 8px;">https://github.com/Mxoder/TinyStories

ingFang SC', 'Hiragino Sans GB', 'Microsoft YaHei UI', 'Microsoft YaHei', Arial, sans-serif;font-size: 1.2em;font-weight: bold;display: table;margin: 2em auto 1em;padding-right: 1em;padding-left: 1em;border-bottom: 2px solid rgb(250, 81, 81);color: rgb(63, 63, 63);">1. 前期准备

ingFang SC', 'Hiragino Sans GB', 'Microsoft YaHei UI', 'Microsoft YaHei', Arial, sans-serif;margin: 1.5em 8px;letter-spacing: 0.1em;color: rgb(63, 63, 63);">让我们先来想一想大概需要做什么。

ingFang SC', 'Hiragino Sans GB', 'Microsoft YaHei UI', 'Microsoft YaHei', Arial, sans-serif;margin: 1.5em 8px;letter-spacing: 0.1em;color: rgb(63, 63, 63);">首先是模型架构的选择。原工作用的是 GPT Neo 架构（可以看他们的 config），这个算是很老的模型了，最初是 EleutherAI 用来复现追踪 GPT-3 的工作的，现在用的也比较少了。我打算选用 LLaMA 架构，也算是符合研究主流、便于推广。LLaMA 3 主要多了个 GQA，也是现在模型的主流，我这里也用一下。

ingFang SC', 'Hiragino Sans GB', 'Microsoft YaHei UI', 'Microsoft YaHei', Arial, sans-serif;margin: 1.5em 8px;letter-spacing: 0.1em;color: rgb(63, 63, 63);">其次是数据的选择。既然是复现，就直接贯彻拿来主义，用原工作开源的数据集（主要是从头生成要花不少 api 费用）。原工作第一版的时候用的是 GPT-3.5 生成的数据，后面社区有人更新了第二版，是用 GPT-4 生成的，比原数据更好，就用它了。

ingFang SC', 'Hiragino Sans GB', 'Microsoft YaHei UI', 'Microsoft YaHei', Arial, sans-serif;margin: 1.5em 8px;letter-spacing: 0.1em;color: rgb(63, 63, 63);">最后是训练。其实我手上就两张 3060 12G 和 4060 Ti 16G，训这个确实是绰绰有余，但我还是不想在桌前吵我自己，于是继续用 Colab。现在 Colab 可以直接看到剩余使用时长了，虽然已经被砍到只有 3h 左右的用卡时间，但至少心里有个底，况且 3h 训我们这个也完全够了。

ingFang SC', 'Hiragino Sans GB', 'Microsoft YaHei UI', 'Microsoft YaHei', Arial, sans-serif;margin: 1.5em 8px;letter-spacing: 0.1em;color: rgb(63, 63, 63);">我们这次用到的 Hugging Face 的库如下：

transformers
accelerate
datasets

理论上比较新的版本都没问题，但如果你很久没更新了，最好用 pip install -U 来升级一下。

我这里的用到的库版本如下，供参考：

torch==2.2.1
transformers==4.40.0
accelerate==0.29.3
datasets==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 机制。

综上，我们的配置如下：

#模型配置
fromtransformersimportAutoConfig

hidden_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(beginofsentence)的id
'eos_token_id':2,#eos_token(endofsentence)的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 中会占用太多的参数比重，并且这只是个专有任务数据训练，没必要用太大的词表。

#分词器
fromtransformersimportAutoTokenizer

tokenizer=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>)
#从右边填充
Onceuponatime</s></s></s></s>...
#从左边填充
</s></s></s></s>Onceuponatime...

3.3 模型实例化

接下来就是实例化模型，这里就不用从预训练模型加载 from_pretrained() 了，而是从配置加载 from_config()：

#模型
importtorch
fromtransformersimportAutoModelForCausalLM

#能用cuda就用cuda
device='cuda'iftorch.cuda.is_available()else'cpu'

#从配置加载模型
model=AutoModelForCausalLM.from_config(
config,
torch_dtype=torch.float32#全精度训练
).to(device)#迁移到device上

'''
LlamaForCausalLM(
(model)lamaModel(
(embed_tokens):Embedding(32000,256)
(layers):ModuleList(
(0-3):4xLlamaDecoderLayer(
(self_attn)lamaSdpaAttention(
(q_proj)inear(in_features=256,out_features=256,bias=False)
(k_proj)inear(in_features=256,out_features=128,bias=False)
(v_proj)inear(in_features=256,out_features=128,bias=False)
(o_proj)inear(in_features=256,out_features=256,bias=False)
(rotary_emb)lamaRotaryEmbedding()
)
(mlp)lamaMLP(
(gate_proj)inear(in_features=256,out_features=768,bias=False)
(up_proj)inear(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 机制。

此时，模型已经初始化了，让我们来打印一下看看参数：

#打印模型的每一层及其参数大小
defprint_model_parameters(model):
print('LayerName&Parameters')
print('----------------------------')
total_params=0
forname,parameterinmodel.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'TotalParameters:{total_params}({total_params/1000000:.1f}M)')

print_model_parameters(model)

得到结果如下：

LayerName&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
----------------------------
TotalParameters:19532032(19.5M)

可以看到，我们的模型只有不到 20M！非常非常小，并且其中 Embedding 占了大头。

尽管模型还没有训练，但我们仍然可以测试一下推理：

definference(
model:AutoModelForCausalLM,
tokenizer:AutoTokenizer,
input_text:str='Onceuponatime,',
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)

'''
Onceuponatime,Hostยcrimeine/\könnenlinewidthmeasurementresolperfectlyTaylormeasèresionesassetviron
'''

嗯，的确是胡言乱语呢，不过可以正常推理，说明模型没问题！

但现在模型是随机初始化的，为了让模型更好地收敛，我们最好给模型一个更好的初始化方法，我这里选用 Kaiming 初始化，比较适用于 ReLU 类的激活，当然也可以选用高斯初始化、Xavier 初始化等等。

#Kaiming初始化
defkaiming_initialization(model):
forname,paraminmodel.named_parameters():
if'weight'innameandparam.dim()>1:
torch.nn.init.kaiming_uniform_(param,mode='fan_in',nonlinearity='leaky_relu')
elif'bias'inname:
#一般偏置项可以初始化为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

#加载数据集
fromdatasetsimportload_dataset

dataset_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':['Onceuponatime,therewasareliableotternamedOllie.Helivedinariverwithhisfamily.Theyalllovedtoplayandswimtogether.\nOneday,Ollie\'smomsaid,'Ollie,hurryandgetsomefishfordinner!'Ollieswamfasttocatchfish.Hesawhisfriend,theduck.'Hi,Ollie!'saidtheduck.'Hi,duck!'saidOllie.'Ineedtohurryandcatchfishformyfamily.'\nWhileOlliewascatchingfish,hefoundabigshinystone.Hethought,'Thisisnotafish,butitissopretty!'Ollietooktheshinystonehometoshowhisfamily.Theyalllookedattheshinystoneandsmiled.Theshinystonemadeeveryonehappy,andtheyforgotaboutthefishfordinner.',
'Oneday,alittleboynamedTimwenttothepark.Hesawabigtiger.Thetigerwasnotmean,butveryeasytoplaywith.Timandthetigerplayedallday.Theyhadlotsoffun.\nThen,somethingunexpectedhappened.Thetigerstartedtoshake.Timwasscared.Hedidnotknowwhatwasgoingon.Butthen,thetigerturnedintoanicedog.Timwasverysurprised.\nTimandthedogplayedtogethernow.Theywereveryhappy.Thedogwaseasytoplaywithtoo.Attheendoftheday,Timwenthomewithhisnewfriend.']}
'''

这里需要注意，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。让我们来写一个处理函数：

fromtypingimportDict,List

defprocess_func(
examplesict[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=[],[]
forinput_idsininput_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 <s>。

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='Runningtokenizerontrain_set:'
)
ds_val=ds_val.map(
process_func,
batched=True,
num_proc=num_proc,
remove_columns=ds_val.column_names,
desc='Runningtokenizeronval_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：

...onceuponatimetherewas...#input
onceuponatimetherewas......#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：

fromtransformersimportDataCollatorForLanguageModeling

data_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)fortextindata]

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]])
}
'''

可以看到：

•由于长度不一，所以 data collator 做了 padding，padding 的方向就是我们 3.2 中提到的 tokenizer.padding_size
•labels 确实是 input_ids 的原位复制，区别在于 input_ids 里用 pad_token_id 来填充，labels 里对应的是 -100、表示不计算 loss

那么微调里常用的 DataCollatorForSeq2Seq 又是如何作用的呢？我们仍然用刚刚的数据例子：

#DataCollatorForSeq2Seq
#这里的tokenizer选用的是Qwen1.5的，并非LLaMA的，只是做一个示意
dc=DataCollatorForSeq2Seq(tokenizer)
data=['南京','南京市','南京市长江']

raw_tokens=[tokenizer(text)fortextindata]

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=[]
fortext,labelindata:
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]])
}
'''

可以看到：

•DataCollatorForSeq2Seq 需要指定当前输入的文本和后面需要生成的文本，即 text 和 label，如果像 DataCollatorForLanguageModeling 那样处理会得不到 labels 字段
•因此，DataCollatorForSeq2Seq 适合有监督微调（SFT），输入是 text，输出是 label，非常合理

5. 超迷你 LLaMA，启动！

5.1 配置训练参数

我们需要用到 transformers 的 TrainingArguments 来配置训练参数，具体参数说明可以看这里。

fromtransformersimportTrainingArguments

training_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=nottorch.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，现在想禁用掉，可以设置环境变量：

importos

os.environ['WANDB_DISABLED']='true'

5.2 配置 Trainer

同样地，具体参数说明可以看这里。

fromtransformersimportTrainer

trainer=Trainer(
model=model,#模型实例
args=training_args,#训练参数
train_dataset=ds_train,#训练集
eval_dataset=ds_val,#验证集（评估集）
tokenizer=tokenizer,#分词器
data_collator=data_collator,#datacollator
)

5.3 训练与保存

配置好 Trainer 后，通过下列代码即可启动训练：

trainer.train()

接下来只需要等待训练完成。我用一个半小时训练了 2 epochs，loss 达到了 1.6 左右。

训练完成后，如果用的是 jupyter notebook，那么此时 model 已经是训练好的状态了。我们可以再次推理试试看：

inference(
model,
tokenizer,
'Onceuponatime,inabeautifulgarden,therelivedalittlerabbitnamedPeterRabbit.',
max_new_tokens=256
)

得到如下结果：

Onceuponatime,inabeautifulgarden,therelivedalittlerabbitnamedPeterRabbit.PeterhadafriendnamedRosie.Theylovedtoplaytogether.Theywouldrun,jump,andlaughalldaylong.
Oneday,Robbysawabigboxinhisyard.Hewascuriousandwantedtoknowwhatwasinside.So,hewenttohisfriend'shouseandasked,'Whatareyoudoing,Spark?'Mayreplied,'Iammakingthisbigboxinthegarden,andIamtryingtoopenit!'
TimmyandHopperwenttofindthebigbox.Theyfoundakeyunderatree.Theyopenedtheboxandfoundmanytoysinside.Theyweresohappytohaveafundaywiththeirnewfriend.Theyplayedwiththetoysalldaylong.Andfromthatdayon,wheneverElliewasapartofsomething,theywouldalwaysrememberthedaytheymetbythebigpond.

可以看到：

•20M 模型确实能够流畅续写故事了
•20M 模型写出的故事的语法、流畅度都不错，但是一致性欠佳，特别是故事主题、人名的前后连贯性不高

总之，这个超迷你 LLaMA 3 确实训练完成了！我们可以将它保存到本地：

model_path='...'

model.save_pretrained(model_path)

也可以推送到 Hugging Face：

fromhuggingface_hubimportnotebook_login

repo_name='TinyStories-LLaMA2-20M-256h-4l-GQA'

notebook_login()#输入AccessTokens

model.push_to_hub(repo_name)
tokenizer.push_to_hub(repo_name)

6. 结尾

这次尝试用 Trainer 来做一个模型的预训练，以往都是用 Trainer 来做微调，这次也算是学习了一下吧。TinyStories 这个工作之前就有关注过，但一直没顾上来复现一下，这次也算是简单复现了个小模型出来，和原工作的丰富度确实是比不了，但也算完成一个 todo。

后面这个小模型可以继续做 SFT，也就是做指令微调，可以和原工作一样，给定故事背景、关键词、开头让小模型续写，也可以迁移到别的任务。不过由于我们的预训练任务只针对了讲短篇故事这一类任务，加上参数又特别少，如果直接迁移其它指令任务估计表现不会很好。

本篇文章也是希望尽可能写得详细一点，欢迎大家交流。

致谢 - Reference

1、LLM大模型之Trainer以及训练参数-知乎(zhihu.com)：详细说明了 Trainer 与 TrainingArguments 的参数，这位博主的其他文章也非常好，值得拜读
https://zhuanlan.zhihu.com/p/662619853
2、NLP（九十四）transformers模块中的DataCollator -知乎(zhihu.com)：详细介绍了 DataCollator，并配上了示例，通俗易懂
https://zhuanlan.zhihu.com/p/692332774
3、stanleylsx/llms_tool:一个基于HuggingFace开发的大语言模型训练、测试工具。支持各模型的webui、终端预测，低参数量及全参数模型训练(预训练、SFT、RM、PPO、DPO)和融合、量化。(github.com)：借鉴了里面采用 DataCollatorForLanguageModeling 来做预训练 data collator 的思路，这个仓库也写得非常工整，值得细读
https://github.com/stanleylsx/llms_tool

欢迎光临链载Ai (https://www.lianzai.com/)