预训练 – GPT2的预训练过程

3 模型加载

import math
import numpy as np
import torch
import torch.nn as nn
from torch.nn import functional as F
import config


# TODO 2 Maybe is ok
# RoPE blogs : https://zhuanlan.zhihu.com/p/645263524
class RoPeEmbedding(nn.Module):
    def __init__(self, config):
        super(RoPeEmbedding, self).__init__()

    def get_freqs_cis(self, d_model, max_seq_len, device):
        # 获取弧度 (d_model/2, )
        freqs = 1. / (10000 ** (torch.arange(0, d_model, 2).float() / d_model))
        # 获取索引 (max_seq_len, )
        t = torch.arange(0, max_seq_len)
        # 得到全索引弧度 (max_seq_len, d_model/2)
        freqs = torch.outer(t, freqs)
        # 位置复数 (max_seq_len, d_model/2)
        freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
        return freqs_cis.to(device)

    def forward(self, xq, xk):
        xq_device = xq.device
        max_seq_len, d_model = xq.shape[1], xq.shape[2]
        # 将xq, xk转为复数
        # (batch_size, max_seq_len, d_model) -> (batch_size, max_seq_len, d_model/2, 2)·
        # -> (batch_size, max_seq_len, d_model/2)
        xq_out = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
        xk_out = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
        # 将xq, xk与旋转复数矩阵进行复数乘法
        freqs_cis = self.get_freqs_cis(d_model, max_seq_len, xq_device)
        # 将最后两个维度变成一维
        xq_out = torch.view_as_real(xq_out * freqs_cis).flatten(-2)
        xk_out = torch.view_as_real(xk_out * freqs_cis).flatten(-2)
        return xq_out.type_as(xq), xk_out.type_as(xk)


# TODO 3 Maybe is ok
class SwishGLU(nn.Module):
    def __init__(self, config):
        super(SwishGLU, self).__init__()
        d_model, dropout = config.d_model, config.dropout
        self.w1 = nn.Linear(d_model, d_model*2)
        self.w3 = nn.Linear(d_model, d_model*2)
        self.w2 = nn.Linear(d_model*2, d_model)
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        x = self.w2(F.silu(self.w1(x)) * self.w3(x))
        x = self.drop(x)
        return x


# TODO 4 Maybe is ok
class RMSNorm(nn.Module):
    def __init__(self, config):
        super(RMSNorm, self).__init__()
        d_model, bias = config.d_model, config.bias
        self.weight = nn.Parameter(torch.ones(d_model))
        self.bias = nn.Parameter(torch.zeros(d_model)) if bias else None
        self.eps = 1e-5

    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)

    def forward(self, x):
        x = self.weight * self._norm(x) + self.bias
        return x


class LayerNorm(nn.Module):
    def __init__(self, config):
        d_model, bias = config.d_model, config.bias
        super(LayerNorm, self).__init__()
        self.weight = nn.Parameter(torch.ones(d_model))
        self.bias = nn.Parameter(torch.zeros(d_model)) if bias else None

    def forward(self, x):
        return F.layer_norm(x, normalized_shape=self.weight.shape, weight=self.weight, bias=self.bias, eps=1e-5)


class SelfAttention(nn.Module):
    def __init__(self, config):
        super(SelfAttention, self).__init__()
        in_channel, out_channel = config.d_model, config.d_model//config.n_head
        self.pe_encoder = config.pe_encoder
        self.rope = RoPeEmbedding(config)
        self.Wq = nn.Linear(in_channel, out_channel)
        self.Wk = nn.Linear(in_channel, out_channel)
        self.Wv = nn.Linear(in_channel, out_channel)

    def forward(self, x):
        q, k, v = self.Wq(x), self.Wk(x), self.Wv(x)
        if self.pe_encoder == "rope":
            q, k = self.rope(q, k)
        attention = q @ k.transpose(-1, -2)
        sqrt_hidden = torch.sqrt(torch.tensor(q.shape[-1], device=q.device))
        mask = torch.tril(torch.ones(size=(q.shape[1], q.shape[1]), device=q.device), diagonal=0)
        attention = (attention / sqrt_hidden).masked_fill(mask == 0, float("-inf"))
        attention = F.softmax(attention, dim=-1)
        value = attention @ v
        return value


class MultiSelfAttention(nn.Module):
    def __init__(self, config):
        super(MultiSelfAttention, self).__init__()
        d_model, n_head = config.d_model, config.n_head
        self.Wo = nn.Linear(d_model, d_model)
        self.multi_self_attention = nn.ModuleList([SelfAttention(config) for _ in range(n_head)])

    def forward(self, x):
        value = self.Wo(torch.cat([self_attention(x) for self_attention in self.multi_self_attention], dim=-1))
        return value


class FFN(nn.Module):
    def __init__(self, config):
        super(FFN, self).__init__()
        d_model, dropout = config.d_model, config.dropout
        self.config = config
        if config.ffn == "origin":
            self.fc1 = nn.Linear(d_model, d_model * 4)
            self.fc2 = nn.Linear(d_model * 4, d_model)
            self.gelu = nn.GELU()
            self.dropout = nn.Dropout(dropout)
        elif config.ffn == "swishglu":
            self.swishglu = SwishGLU(config)

    def forward(self, x):
        if self.config.ffn == "origin":
            x = self.gelu(self.fc1(x))
            x = self.dropout(self.fc2(x))
        elif self.config.ffn == "swishglu":
            x = self.swishglu(x)
        return x


class TransformerLayer(nn.Module):
    def __init__(self, config):
        super(TransformerLayer, self).__init__()
        if config.norm == "origin":
            self.norm1 = LayerNorm(config)
            self.norm2 = LayerNorm(config)
        elif config.norm == "rms":
            self.norm1 = RMSNorm(config)
            self.norm2 = RMSNorm(config)
        self.multi_self_attention = MultiSelfAttention(config)
        self.ffn = FFN(config)


    def forward(self, x):
        x = x + self.multi_self_attention(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x


class GPT2(nn.Module):
    def __init__(self, config):
        super(GPT2, self).__init__()
        self.config = config
        self.transformer = nn.ModuleDict(dict(
            vocab = nn.Embedding(config.vocab_size, config.d_model),
            pe = nn.Embedding(config.max_seq_len, config.d_model) if config.pe_encoder == "origin" else None,
            drop = nn.Dropout(config.dropout),
            h = nn.ModuleList([TransformerLayer(config) for _ in range(config.n_layer)]),
            norm = LayerNorm(config) if config.norm == "origin" else RMSNorm(config)
        ))
        self.lm_head = nn.Linear(config.d_model, config.vocab_size)
        self.apply(self._init_weights)
        for name, paramters in self.transformer.named_parameters():
            if name.endswith("fc2.weight"):
                torch.nn.init.normal_(paramters, 0, (0.02 / math.sqrt(2 * config.n_layer)))

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, 0, 0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        if isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, 0, 0.02)

    def forward(self, x, target=None):
        token_emb = self.transformer.vocab(x)
        if self.config.pe_encoder == "origin":
            pos = torch.arange(0, x.shape[1], dtype=torch.long, device=x.device)
            pos_emb = self.transformer.pe(pos)
            h = self.transformer.drop(token_emb + pos_emb)
        else:
            h = self.transformer.drop(token_emb)
        for block in self.transformer.h:
            h = block(h)
        h = self.transformer.norm(h)
        if target is not None:
            logits = self.lm_head(h)
            loss = F.cross_entropy(logits.view(-1, logits.shape[-1]), target.view(-1), ignore_index=-1)
        else:
            logits = self.lm_head(h[:, -1, :])
            loss = None
        return logits, loss

    def generate(self, idx, max_tokens, t=0.1, topk=None):
        for _ in range(max_tokens):
            logits = self(idx)
            logit = logits[:, -1, :] / t
            if topk is not None:
                # vector index
                v, _ = torch.topk(logit, min(topk, logit.shape[-1]))
                logit[logit < v[:, [-1]]] = -float('inf')
            probs = F.softmax(logit, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=-1)
        return idx

# TODO 1 增加推理代码 K/V cache 参考LLaMA
# TODO 2 增加RoPE 参考LLaMa  OK
# TODO 3 更换SwishGLU OK
# TODO 4 RMSNorm     OK
# TODO 5 增加其他2种采样方法, top_p, beam_search
# TODO 6 FlashAttention 参考LLaMA-moe



if __name__ == "__main__":
    config = config.config_args
    gpt2 = GPT2(config).to("cuda")
    num_parameters = np.sum([p.numel() for p in gpt2.parameters()]) / 1024/ 1024
    print(f"paramters is {num_parameters} M")
    x = torch.randint(0, 1000, size=(2, 10), dtype=torch.long).to("cuda")
    print(f"x is {x}")
    # output = gpt2.generate(x, config.max_tokens, t=0.1, topk=5)
    # print(f"output is {output}")
    logits, loss = gpt2(x)
    print(logits)