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/TinyStoriesingFang 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 的库如下：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 机制。 综上，我们的配置如下： #模型配置fromtransformersimportAutoConfighidden_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 中会占用太多的参数比重，并且这只是个专有任务数据训练，没必要用太大的词表。 #分词器fromtransformersimportAutoTokenizertokenizer=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()： #模型importtorchfromtransformersimportAutoModelForCausalLM#能用cuda就用cudadevice='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=0forname,parameterinmodel.named_parameters():param_size=parameter.size()param_count=torch.prod(torch.tensor(param_size)).item()total_params+=param_countprint(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:8192000model.layers.0.self_attn.q_proj.weight|Size:torch.Size([256,256])|Count:65536model.layers.0.self_attn.k_proj.weight|Size:torch.Size([128,256])|Count:32768model.layers.0.self_attn.v_proj.weight|Size:torch.Size([128,256])|Count:32768model.layers.0.self_attn.o_proj.weight|Size:torch.Size([256,256])|Count:65536model.layers.0.mlp.gate_proj.weight|Size:torch.Size([768,256])|Count:196608model.layers.0.mlp.up_proj.weight|Size:torch.Size([768,256])|Count:196608model.layers.0.mlp.down_proj.weight|Size:torch.Size([256,768])|Count:196608中间省略...model.layers.3.input_layernorm.weight|Size:torch.Size([256])|Count:256model.layers.3.post_attention_layernorm.weight|Size:torch.Size([256])|Count:256model.norm.weight|Size:torch.Size([256])|Count:256lm_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:#一般偏置项可以初始化为0torch.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_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':['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,Listdefprocess_func(examples ict[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 数据批处理——DataCollator4.3.1 两行代码我们在训练的时候往往不会一条一条训练，而是成批次地训练，那么我们就需要对数据做批处理。因此我们需要用到 transformers 中的一个工具系列——DataCollator[1]。 既然是预训练，那么就是让模型在语料上做无监督学习，也就是我们熟知的 next token prediction，即根据前面的所有输入来预测下一个 token，然后把新的 token 拼接在已有输入上作为下一输入，如此往复，直到触发停止设定（例如触发 max_new_tokens）。 所以我们的训练目标——或者说是 label——显而易见，就是把输入偏移一位当作预测目标，我们计算的就是输出和这个目标之间的 loss： ...onceuponatimetherewas...#inputonceuponatimetherewas......#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： fromtransformersimportDataCollatorForLanguageModelingdata_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]])}'''可以看到： 那么微调里常用的 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]])}'''可以看到： 5. 超迷你 LLaMA，启动！5.1 配置训练参数我们需要用到 transformers 的 TrainingArguments 来配置训练参数，具体参数说明可以看这里。 fromtransformersimportTrainingArgumentstraining_args=TrainingArguments(output_dir='saves',#输出路径，包括模型检查点、中间文件等overwrite_output_dir=True,#是否覆写output_dirdo_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(),#尝试配置bf16fp16=nottorch.cuda.is_bf16_supported(),#bf16不行就上fp16logging_steps=50,#打印步骤间隔report_to=None,#日志输出目标，不想用wandb可以设置为Nonenum_train_epochs=2,#训练轮数，2~3即可save_steps=1000,#检查点保存步骤间隔save_total_limit=2,#output_dir内留存的检查点最大数目seed=3407#随机种子) 如果你之前用了 wandb，现在想禁用掉，可以设置环境变量： importosos.environ['WANDB_DISABLED']='true' 5.2 配置 Trainer同样地，具体参数说明可以看这里。 fromtransformersimportTrainertrainer=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. 可以看到： 总之，这个超迷你 LLaMA 3 确实训练完成了！我们可以将它保存到本地： model_path='...'model.save_pretrained(model_path) 也可以推送到 Hugging Face： fromhuggingface_hubimportnotebook_loginrepo_name='TinyStories-LLaMA2-20M-256h-4l-GQA'notebook_login()#输入AccessTokensmodel.push_to_hub(repo_name)tokenizer.push_to_hub(repo_name) 6. 结尾这次尝试用 Trainer 来做一个模型的预训练，以往都是用 Trainer 来做微调，这次也算是学习了一下吧。TinyStories 这个工作之前就有关注过，但一直没顾上来复现一下，这次也算是简单复现了个小模型出来，和原工作的丰富度确实是比不了，但也算完成一个 todo。 后面这个小模型可以继续做 SFT，也就是做指令微调，可以和原工作一样，给定故事背景、关键词、开头让小模型续写，也可以迁移到别的任务。不过由于我们的预训练任务只针对了讲短篇故事这一类任务，加上参数又特别少，如果直接迁移其它指令任务估计表现不会很好。
|
