Transformers的底层torch的认识（一）

transformers

我们经常使用transformers库，使用它加载模型加载模型权重、加载分词器、保存模型权重等等但实际上，transformers 库的内部模型加载类使用了 PyTorch（以及一些场景下的 TensorFlow）来实现大部分功能，特别是模型的加载、保存和推理。在 PyTorch 中，模型的加载和保存主要依赖 torch 提供的工具函数，如 torch.load、torch.save、state_dict 等。

python
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
    # 加载模型权重到cpu
    state_dict = torch.load(resolved_archive_file, map_location='cpu')
    # 将加载的权重设置到模型中
    model.load_state_dict(state_dict)
    
    
    
def save_pretrained(self, save_directory, **kwargs):
    # 将模型的状态字典(即权重)保存到指定目录
    torch.save(self.state_dict(), os.path.join(save_directory, WEIGHTS_NAME))
    # 保存模型配置
    self.config.save_pretrained(save_directory)




def resize_token_embeddings(self, new_num_tokens):
    # 获取旧的词嵌入层
    old_embeddings = self.get_input_embeddings()
    # 创建新的嵌入层，其大小为 new_num_tokens
    new_embeddings = torch.nn.Embedding(new_num_tokens, old_embeddings.weight.size(1))
    # 将旧的嵌入权重复制到新的嵌入层中
    new_embeddings.weight.data[:old_embeddings.num_embeddings] = old_embeddings.weight.data
    # 将新的嵌入层设置为模型的输入嵌入层
    self.set_input_embeddings(new_embeddings)
    
    
    
    
def _init_weights(self, module):
    # 初始化线性层的权重
    if isinstance(module, torch.nn.Linear):
        torch.nn.init.xavier_uniform_(module.weight)
    # 初始化嵌入层的权重
    elif isinstance(module, torch.nn.Embedding):
        torch.nn.init.normal_(module.weight, mean=0, std=self.config.initializer_range)

transformers你不能很好的了解到模型底层的算法，所以学习一下torch在语言模型方面的代码与实践！

构建语言模型常用torch模块

嵌入层 (nn.Embedding): 将词汇映射到高维空间中的密集向量。
循环层:
- nn.RNN: 基本的循环神经网络层transformers 库的内部模型加载类使用了 PyTorch（以及一些场景下的 TensorFlow）来实现大部分功能，特别是模型的加载、保存和推理。在 PyTorch 中，模型的加载和保存主要依赖 torch 提供的工具函数，如 torch.load、torch.save、state_dict 等
- nn.Transformer: 完整的 Transformer 模型，包括自注意力机制和前馈网络。
- nn.TransformerEncoder: Transformer 编码器部分。
- nn.TransformerDecoder: Transformer 解码器部分。
注意力机制:
- nn.MultiheadAttention: 多头自注意力机制，是 Transformer 架构的核心部分。
位置编码:
- nn.Embedding: 虽然不是专门的位置编码模块，但通常使用嵌入层来实现位置编码。
激活函数:
- nn.ReLU, nn.GELU, nn.SELU 等：在 Transformer 的前馈网络中使用。
层归一化 (nn.LayerNorm): 用于归一化层的输出，有助于加速训练并提高模型稳定性。
Dropout层 (nn.Dropout): 用于在训练过程中随机丢弃神经元，以减少过拟合。
线性层 (nn.Linear): 用于实现全连接层，常见于 Transformer 的前馈网络和解码器的输出层。
损失函数:
- nn.CrossEntropyLoss: 用于语言模型的分类任务，如词预测。
- nn.BCEWithLogitsLoss: 用于二分类任务，如情感分析。
优化器:
- nn.Adam, nn.SGD 等：用于模型的训练。
其他辅助模块:
- nn.ModuleList: 用于存放有序的模块列表。
- nn.Parameter: 用于定义自定义参数。

llama 8b instruct

LlamaForCausalLM
- (model): LlamaModel
  - (embed_tokens): Embedding(128256, 4096)
  - (layers): ModuleList
    - (0-31): 32 x LlamaDecoderLayer
      - (self_attn): LlamaSdpaAttention
        
        (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
        
        (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
        
        (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
        
        (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
        
        (rotary_emb): LlamaRotaryEmbedding()
      - (mlp): LlamaMLP
        
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        
        (act_fn): SiLU()
      - (input_layernorm): LlamaRMSNorm((4096,), eps=1e-05)
      - (post_attention_layernorm): LlamaRMSNorm((4096,), eps=1e-05)
  - (norm): LlamaRMSNorm((4096,), eps=1e-05)
  - (rotary_emb): LlamaRotaryEmbedding()
- (lm_head): Linear(in_features=4096, out_features=128256, bias=False)

多层感知机的简单实现

python
import torch
from torch import nn
from d2l import torch as d2l

python
net = nn.Sequential(nn.Flatten(),
                    nn.Linear(784, 256),
                    nn.ReLU(),
                    nn.Linear(256, 10))

def init_weights(m):
    if type(m) == nn.Linear:
        nn.init.normal_(m.weight, std=0.01)

net.apply(init_weights);

python
batch_size, lr, num_epochs = 256, 0.1, 10
loss = nn.CrossEntropyLoss(reduction='none')
trainer = torch.optim.SGD(net.parameters(), lr=lr)

train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
d2l.train_ch3(net, train_iter, test_iter, loss, num_epochs, trainer)