从头预训练一只超迷你 LLaMA 3

https://github.com/Mxoder/TinyStories

transformers
accelerate
datasets

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

# 模型配置
from transformers import AutoConfig

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 (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
}
'''

# 分词器
from transformers import AutoTokenizer

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

# 假设 pad_token 就是 eos_token(</s>)
# 从右边填充
Once upon a time </s></s></s></s>...
# 从左边填充
</s></s></s></s>Once upon a time ...

# 模型
import torch
from transformers import AutoModelForCausalLM

# 能用 cuda 就用 cuda
device = '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)
)
'''

# 打印模型的每一层及其参数大小
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)

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 初始化
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)

https://huggingface.co/datasets/noanabeshima/TinyStoriesV2

# 加载数据集
from datasets import load_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': ['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.']}
'''

# 数据长这样
{
    'text': [
        <text_1>, <text_2>, <text_3>, ...
    ]
}

# 不是下面这样
[
    {'text': <text_1>},
    {'text': <text_2>},
    {'text': <text_3>},
    ...
]

from typing import Dict, List

def 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
    }

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

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))

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

...   once  upon  a     time   there  was  ...        # input
once  upon  a     time  there  was    ...  ...        # label

# label 错开一位，是 input 的下一预测目标，计算的就是 input 和 label 之间的 loss

from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

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

# 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
# 这里的 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
# 这里的 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]])
}
'''

from transformers import TrainingArguments

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=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                                   # 随机种子
)

import os

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

from transformers import Trainer

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

trainer.train()

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.

model_path = '...'

model.save_pretrained(model_path)

from huggingface_hub import notebook_login

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

notebook_login()    # 输入 Access Tokens

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

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