Llama系模型總結

幻影星全能的木豆發表於2024-07-19

Llama3 學習連結 https://blog.csdn.net/v_JULY_v/article/details/137955982

就不易理解的內容進一步剖析

對Llama系模型進行彙總

目錄
  • 一、LLama1
    • 1. LLama 1 簡介
    • 2. 模型架構
      • 2.1 RMSNorm
      • 2.2 SwiGLU替代ReLU
      • 2.3 位置編碼:RoPE
      • 2.4 Transform架構的實現:Attention計算、SA、FFN
        • SA
        • FFN
        • Transformer Block
        • Transformer Decoder
        • 生成過程
      • 2.5 LLaMA的Optimizer設計、模型加速最佳化與微型版本
  • 二、LLama2
    • 1. 簡介
    • 2. 模型架構
      • 2.1 LLaMA 2-Chat:三個版本——7B 13B 70B
      • 2.2 分組查詢注意力——Grouped-Query Attention
    • 3. Llama 2-Chat中的RLHF
      • 3.1 監督微調(SFT)
      • 3.2 訓練兩個獎勵模型:一個偏實用 一個偏安全
      • 2.6 具體的策略迭代:PPO與拒絕取樣
  • 三、LLama3
    • 1. Llama 3 簡介
    • 2. 模型架構
      • 2.1 擴充詞表
      • 2.2 採用分組查詢注意力 GQA
      • 2.3 上下文擴充套件至8K
    • 3. 訓練資料
    • 4. 提高預訓練效率
    • 5. 後訓練策略
    • 6. Llama 3 的下一步
  • 四、擴充套件Llama3至100萬上下文
    • 1. Llama-3-8B 的上下文長度擴充套件到16K
    • 2. LLama3長度擴充套件到100萬

一、LLama1

1. LLama 1 簡介

23年2.24日, Meta釋出了LLaMA, 有多個引數規模的版本(7B 13B 33B 65B)

LLaMA只使用公開的資料(總計1.4T即1,400GB的token,其中CommonCrawl的資料佔比67%,C4資料佔比15%,Github、Wikipedia、Books這三項資料均都各自佔比4.5%,ArXiv佔比2.5%,StackExchange佔比2%)

證明小模型在足夠多的的資料上訓練後,也能達到甚至超過大模型的效果

比如13B引數的版本在多項基準上測試的效果好於2020年的引數規模達175B的GPT-3

而對於65B引數的LLaMA,則可與DeepMind的Chinchilla(70B引數)和谷歌的PaLM(540B引數)旗鼓相當

2. 模型架構

2.1 RMSNorm

為了提高訓練的穩定性,對每個transformer子層的輸入進行歸一化,而不是對輸出進行歸一化,且使用由Zhang和Sennrich(2019)提出的RMSNorm

RMS Norm是一般LayerNorm的一種變體,可以在梯度下降時令損失更加平滑

與layerNorm相比,RMS Norm的主要區別在於去掉了減去均值的部分(re-centering),只保留方差部分(re-scaling)

為一目瞭然,我們看下它們各自的歸一化的表示式

LayerNorm

在給定一個輸入特徵向量 \(a_i\)後,先計算 \(a\) 的均值 μ 和標準差 σ

\(\begin{array}{c} \mu=\frac{1}{n} \sum_{i=1}^{n} a_{i} \\ \sigma=\sqrt{\frac{1}{n} \sum_{i=1}^{n}\left(a_{i}-\mu\right)^{2}} \end{array}\)

然後進行歸一化操作: \(\bar{a}_{i}=\frac{a_{i}-\mu}{\sigma} g_{i}+b_{i}\)

其中的 \(g_i\) 是可學習的縮放引數,來調整每個特徵在歸一化後的尺度或權重,最終作用是恢復歸一化操作可能損失的資訊,如資料的比例和分佈等

\(b_i\) 是偏移因子,可以對歸一化並放縮後的資料進行偏移,使模型可以學習到一個最優的數值範圍,比如在ReLU啟用函式中,我們可能希望值在0以上

RMS Norm

首先,計算輸入特徵向量 a 的平方根均值 \(R M S(a)=\sqrt{\frac{1}{n} \sum_{i=1}^{n} a_{i}{ }^{2}}\)

然後,對輸入特徵向量 a 進行歸一化 \(\bar{a}_{i}=\frac{a_{i}}{R M S(a)} g_{i}\)

此外,可選地,RMSNorm 還可以引入可學習的偏移引數 \(\bar{a}_{i}=\frac{a_{i}}{R M S(a)} g_{i}+b_{i}\)

class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        # eps防止取倒數之後分母為0
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
 
	def _norm(self, x):
		return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
 
    def forward(self, x):
        output = self._norm(x.float()).type_as(x)
        # weight是末尾乘的可訓練引數,即gi
        return output * self.weight

2.2 SwiGLU替代ReLU

為了更好的理解SwiGLU,首先你得先了解什麼是ReLU和GLU

  1. ReLU的函式表示式為 \(f(x)=\max (0, x)\),這意味著對於所有負的輸入值,ReLU函式的輸出都是0,對於所有正的輸入值,ReLU函式的輸出等於輸入值本身。

  2. GLU 的基本思想是引入一種稱為“門”機制,該機制可以動態地控制資訊的流動

    \(G L U(x)=x \otimes \sigma(g(x))\)

    這個公式意味著,對於每個輸入 \(x\),都會有一個相應的門值,這個門值由 \(\sigma(g(x))\)​​ 產生,其範圍在 0 到 1 之間(在正數區域接近於1,負數區域接近於0),這個門值用於調節相應的輸入值

    如果 \(\sigma(g(x))\)​ 接近 1,那麼“門”就幾乎完全開啟,輸入 x 的資訊能夠自由流動,於是 GLU 的輸出接近於 x

    如果 \(\sigma(g(x))\) 接近 0,意味著“門”幾乎完全關閉,即輸入 x 的大部分或全部資訊被阻止透過,於是 GLU 的輸出接近 0

    sigmoid函式 \(\sigma(x)=\frac{1}{1+e^{-x}}\)

    image-20240719180807331

    sigmoid導函式影像

    image-20240719180904618

而LLaMA採用Shazeer(2020)提出的SwiGLU替換了原有的ReLU,SwiGLU的作用機制是根據輸入資料的特性,透過學習到的引數自動調整資訊流動的路徑,具體是採用SwiGLU的Feedforward Neural Network (簡稱FNN,這是一種使用可學習的門控機制的前饋神經網路)

其在論文中以如下公式進行表述:

\(F F N_{swiGLU}\left(x, W, V, W_{2}\right)=\left(\operatorname{Swish}_{\beta}(x W) \otimes x V\right) W_{2}\)

  1. 先是透過Swish非線性啟用函式處理 “輸入x和權重矩陣W的乘積”

    Swish啟用函式 \(f(x)=x * \operatorname{sigmoid}(\beta x)\),輸入被縮放了 \(\beta\) 倍,\(\beta\) 是一個可以學習的引數,比如下圖, \(\beta\) 不同,Swish啟用函式的形狀則各異

    image-20240719181658533

    \(\beta\)​ 趨近於 0 時,Swish 函式趨近於線性函式 y = x

    \(\beta\) 趨近於無窮大時,Swish 函式趨近於 ReLU 函式

  2. 步驟1得到的結果和 “輸入x與權重矩陣V的乘積” 進行逐元素的乘法

    這個操作相當於在 Swish 啟用的輸出和第二個線性變換的輸出(指的是後面的\(W_2\))之間引入了一個類似於GLU的“門”,這個門的值是由原始輸入 x透過線性變換 V計算得到的,因此,它可以動態地控制 Swish 啟用的輸出

  3. 最後乘以權重矩陣\(W_2\)

2.3 位置編碼:RoPE

LLaMA實現的旋轉位置編碼

# 預計算頻率和複數的函式
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))    # 計算頻率
    t = torch.arange(end, device=freqs.device)    # 根據結束位置生成序列
    freqs = torch.outer(t, freqs).float()    # 計算外積得到新的頻率
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)    # 計算複數
    return freqs_cis    # 返回複數
# 重塑的函式
def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
    ndim = x.ndim    # 獲取輸入張量的維度
    assert 0 <= 1 < ndim    # 檢查維度的合理性
    assert freqs_cis.shape == (x.shape[1], x.shape[-1])    # 檢查複數的形狀
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]    # 計算新的形狀
    return freqs_cis.view(*shape)    # 重塑複數的形狀並返回
# 應用旋轉嵌入的函式
def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))    # 將xq視為複數
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))    # 將xk視為複數
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)    # 重塑複數的形狀
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)    # 計算xq的輸出
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)    # 計算xk的輸出
    return xq_out.type_as(xq), xk_out.type_as(xk)    # 返回xq和xk的輸出
# 對Query和Key應用旋轉嵌入
xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

2.4 Transform架構的實現:Attention計算、SA、FFN

SA

Attention計算的總體過程

1.) 輸入\(x\),分別經過三個Linear得到\(x_q, x_k, x_v\)

2.) 在\(x_q, x_k\)中加入旋轉位置編碼

3.) 快取 \(x_q, x_k\)

4.) 計算\(softmax(\frac{QK^T}{\sqrt{d_k}})V\)

class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        # 設定本地注意力頭的數量
        self.n_local_heads = args.n_heads // fs_init.get_model_parallel_world_size()
        # 每個注意力頭的維度
        self.head_dim = args.dim // args.n_heads

        # Query投影層
        self.wq = ColumnParallelLinear(
            args.dim,
            args.n_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        # Key投影層
        self.wk = ColumnParallelLinear(
            args.dim,
            args.n_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        # Value投影層
        self.wv = ColumnParallelLinear(
            args.dim,
            args.n_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        # 輸出投影層
        self.wo = RowParallelLinear(
            args.n_heads * self.head_dim,
            args.dim,
            bias=False,
            input_is_parallel=True,
            init_method=lambda x: x,
        )

        # 使用零初始化鍵快取
        self.cache_k = torch.zeros(
            (args.max_batch_size, args.max_seq_len, self.n_local_heads, self.head_dim)
        ).cuda()
        # 使用零初始化值快取
        self.cache_v = torch.zeros(
            (args.max_batch_size, args.max_seq_len, self.n_local_heads, self.head_dim)
        ).cuda()
 
	def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
        bsz, seqlen, _ = x.shape
        # 進行Query投影
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

        # 將形狀調整為[bsz, seqlen, n_local_heads, head_dim]
        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_local_heads, self.head_dim)

        # 對Query和Key應用旋轉嵌入
        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

        # 將快取鍵和值轉換為xq的裝置型別
        self.cache_k = self.cache_k.to(xq)
        self.cache_v = self.cache_v.to(xq)

        # 更新快取鍵和值
        self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
        self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv

        # 獲取鍵和值
        keys = self.cache_k[:bsz, : start_pos + seqlen]
        values = self.cache_v[:bsz, : start_pos + seqlen]

        # 轉置xq、鍵和值的維度
        xq = xq.transpose(1, 2)
        keys = keys.transpose(1, 2)
        values = values.transpose(1, 2)

        # 計算注意力分數
        scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
        if mask is not None:
        scores = scores + mask  # (bs, n_local_heads, slen, cache_len + slen)
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)

        # 使用注意力分數加權求和得到輸出
        output = torch.matmul(scores, values)  # (bs, n_local_heads, slen, head_dim)
        output = output.transpose(
        1, 2
        ).contiguous().view(bsz, seqlen, -1)

        # 應用輸出投影
        return self.wo(output)
FFN

前饋網路FFN部分,需要注意的點就是採用的啟用函式,以及啟用函式的位置

import torch.nn as nn
import torch.nn.functional as F
 
class FeedForward(nn.Module):
    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        multiple_of: int,
    ):
        super().__init__()
 
        # 初始化隱藏層的維度為輸入維度的2/3
        hidden_dim = int(2 * hidden_dim / 3)
        # 調整隱藏層維度為multiple_of的倍數
        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

        # 第一個線性層
        self.w1 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )

        # 第二個線性層
        self.w2 = RowParallelLinear(
            hidden_dim, dim, bias=False, input_is_parallel=True, init_method=lambda x: x
        )

        # 第三個線性層
        self.w3 = ColumnParallelLinear(
            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
        )

    def forward(self, x):
        # 前向傳播函式
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

與常見模型中的FFN做一下簡單的對比

  • BART中的FFN,用的是fc->act->fc,用了兩層全連線

  • GPT中的FFN,用的是conv1D->act->conv1D,也是隻用了兩層

  • 而LLaMA中的FFN採用了三個全連線層以實現FFNSwiGLU

    \(F F N_{\text {swiGLU }}\left(x, W, V, W_{2}\right)=\left(\operatorname{Swish}_{1}(x W) \otimes x V\right) W_{2}\)

Transformer Block

將SA和FFN這兩部分拼在一起就是一個transformer block

import torch
import torch.nn as nn
from typing import Optional
 
class TransformerBlock(nn.Module):
    def __init__(self, layer_id: int, args: ModelArgs):
        super().__init__()
        # 初始化引數
        self.n_heads = args.n_heads  # 注意力頭的數量
        self.dim = args.dim  # 模型維度
        self.head_dim = args.dim // args.n_heads  # 每個注意力頭的維度
        self.attention = Attention(args)  # 注意力機制模組
        self.feed_forward = FeedForward(
            dim=args.dim, hidden_dim=4 * args.dim, multiple_of=args.multiple_of
        )  # 前饋神經網路模組
        self.layer_id = layer_id  # 當前層的ID
        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)  # 注意力模組的歸一化
        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)  # 前饋神經網路模組的歸一化
        
    def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
        # 輸入x經過self-attention之後,做Add&Norm
        h = x + self.attention.forward(self.attention_norm(x), start_pos, freqs_cis, mask)
        # 上一步的輸出h作為輸入,經過前饋神經網路Feed forward之後,做Add&Norm
        out = h + self.feed_forward.forward(self.ffn_norm(h))
        return out

注意這裡有一個SA的RMSNorm,有一個FFN的RMSNorm

先對x做RMSNorm再進入Self-Attention

先對h做RMSNorm再進入FFN

Transformer Decoder

最後利用torch的 ModuleList 將transformer block進行堆疊,拼上最前頭的embedding部分,就是一個完整的transformer decoder結構了

import torch
import torch.nn as nn
from typing import Optional
 
class Transformer(nn.Module):
    def __init__(self, params: ModelArgs):
        super().__init__()
    # 初始化引數
    self.params = params
    self.vocab_size = params.vocab_size  # 詞彙表大小
    self.n_layers = params.n_layers  # Transformer模型的層數
    # 詞嵌入層
    self.tok_embeddings = ParallelEmbedding(
        params.vocab_size, params.dim, init_method=lambda x: x
    )
    # Transformer的各個層
    self.layers = torch.nn.ModuleList()
    for layer_id in range(params.n_layers):
        self.layers.append(TransformerBlock(layer_id, params))
        
    # 歸一化層
    self.norm = RMSNorm(params.dim, eps=params.norm_eps)
    
    # 輸出層
    self.output = ColumnParallelLinear(
        params.dim, params.vocab_size, bias=False, init_method=lambda x: x
    )
    # 預計算的頻率矩陣
    self.freqs_cis = precompute_freqs_cis(
        self.params.dim // self.params.n_heads, self.params.max_seq_len * 2
    )
    
	@torch.inference_mode()
	def forward(self, tokens: torch.Tensor, start_pos: int):
        _bsz, seqlen = tokens.shape
        # Token嵌入和位置編碼
        h = self.tok_embeddings(tokens)
        self.freqs_cis = self.freqs_cis.to(h.device)
        freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]

        # 生成上三角的mask矩陣(為decoder模型防止標籤洩漏)
        '''
        [0., -∞, -∞, -∞, -∞, -∞, -∞, -∞, -∞, -∞],
        [0., 0., -∞, -∞, -∞, -∞, -∞, -∞, -∞, -∞],
        [0., 0., 0., -∞, -∞, -∞, -∞, -∞, -∞, -∞],
        [0., 0., 0., 0., -∞, -∞, -∞, -∞, -∞, -∞],
        [0., 0., 0., 0., 0., -∞, -∞, -∞, -∞, -∞],
        [0., 0., 0., 0., 0., 0., -∞, -∞, -∞, -∞],
        [0., 0., 0., 0., 0., 0., 0., -∞, -∞, -∞],
        [0., 0., 0., 0., 0., 0., 0., 0., -∞, -∞],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., -∞],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]
        '''
        mask = None
        if seqlen > 1:
            mask = torch.full((1, 1, seqlen, seqlen), float("-inf"), device=tokens.device)
            mask = torch.triu(mask, diagonal=start_pos + 1).type_as(h)

        # 逐層計算Transformer
        for layer in self.layers:
            h = layer(h, start_pos, freqs_cis, mask)
        h = self.norm(h)
        output = self.output(h[:, -1, :])  # 只計算最後一個位置的logits
        return output.float()
生成過程
  1. 對prompts進行tokenize,得到token ids;
  2. 計算當前batch的最大長度total_len,用來建立輸入的token tensor,最大長度不能超過前文所述快取的大小;
  3. 從當前batch中,把最短的一個prompt的位置,作為生成的開始位置,開始生成
  4. 輸入的token tensor傳入transformer模型,計算logits,得到形狀為(batch_size, hidden_size)的logits(transformer最後一層的輸出);
  5. softmax+top_p取樣,得到當前預測的token,並更新當前位置,準備預測下一個token;
  6. 解碼得到生成的文字
class LLaMA:
    def __init__(self, model: Transformer, tokenizer: Tokenizer):
        self.model = model
        self.tokenizer = tokenizer

    def generate(
        self,
        prompts: List[str],
        max_gen_len: int,
        temperature: float = 0.8,
        top_p: float = 0.95,
    ) -> List[str]:
        # 獲取批處理大小
        bsz = len(prompts)
        # 獲取模型引數
        params = self.model.params
        # 檢查批處理大小是否在允許的最大批處理大小範圍內
        assert bsz <= params.max_batch_size, (bsz, params.max_batch_size)

        # 使用分詞器對提示進行編碼為標記
        prompt_tokens = [self.tokenizer.encode(x, bos=True, eos=False) for x in prompts]

        # 查詢提示標記的最小和最大大小
        min_prompt_size = min([len(t) for t in prompt_tokens])
        max_prompt_size = max([len(t) for t in prompt_tokens])

        # 計算要生成的標記的總長度
        total_len = min(params.max_seq_len, max_gen_len + max_prompt_size)

        # 建立一個張量來儲存生成的標記,填充為填充標記
        tokens = torch.full((bsz, total_len), self.tokenizer.pad_id).cuda().long()
        # 將提示標記複製到標記張量中
        for k, t in enumerate(prompt_tokens):
            tokens[k, : len(t)] = torch.tensor(t).long()
        # 建立一個掩碼以識別輸入文字
        input_text_mask = tokens != self.tokenizer.pad_id
        # 設定生成的起始位置
        start_pos = min_prompt_size
        prev_pos = 0
        # 逐個生成標記
        for cur_pos in range(start_pos, total_len):
            # 透過模型進行前向傳遞以獲取logits
            logits = self.model.forward(tokens[:, prev_pos:cur_pos], prev_pos)
            if temperature > 0:
                # 對logits應用溫度並計算機率
                probs = torch.softmax(logits / temperature, dim=-1)
                # 使用top-p取樣抽樣下一個標記
                next_token = sample_top_p(probs, top_p)
            else:
                # 選擇機率最高的標記
                next_token = torch.argmax(logits, dim=-1)
            next_token = next_token.reshape(-1)
            # 只有在已經生成了提示的情況下才替換標記
            next_token = torch.where(
                input_text_mask[:, cur_pos], tokens[:, cur_pos], next_token
            )
            tokens[:, cur_pos] = next_token
            prev_pos = cur_pos

        # 將生成的標記解碼為文字
        decoded = []
        for i, t in enumerate(tokens.tolist()):
            # 將標記截斷到最大生成長度
            t = t[: len(prompt_tokens[i]) + max_gen_len]
            # 將標記截斷到如果存在結束標記
            try:
                t = t[: t.index(self.tokenizer.eos_id)]
            except ValueError:
                pass
            # 將標記解碼為文字
            decoded.append(self.tokenizer.decode(t))
        return decoded

def sample_top_p(probs, p):
    # 按降序對機率進行排序
    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
    # 計算機率的累積和
    probs_sum = torch.cumsum(probs_sort, dim=-1)
    # 建立一個掩碼以過濾累積機率超過p的標記
    mask = probs_sum - probs_sort > p
    # 將被過濾的標記的機率設定為0
    probs_sort[mask] = 0.0
    # 歸一化機率
    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
    # 使用修改後的機率進行抽樣下一個標記
    next_token = torch.multinomial(probs_sort, num_samples=1)
    # 收集抽樣標記的原始索引
    next_token = torch.gather(probs_idx, -1, next_token)
    return next_token

2.5 LLaMA的Optimizer設計、模型加速最佳化與微型版本

在Optimizer設計上

Llama使用AdamW最佳化器(Loshchilov和Hutter,2017)進行訓練,超引數設定為β1=0.9,β2=0.95
此外,使用餘弦學習率方式,使最終學習率等於最大學習率的10%,以及使用0.1的權重衰減和1.0的梯度剪裁,和2000個warm up策略,使得可以根據模型的大小改變學習率和批次大小

在模型的加速最佳化方面

  1. 首先,使用一個高效的因果多頭注意力方式的實現,靈感來自Rabe和Staats(2021)以及Dao等人(2022),這個實現可在xformers庫中找到,可以有效減少記憶體的使用和計算

    具體原理為透過不儲存注意力權重和不計算由於語言建模任務的因果性質而被掩蓋的鍵/查詢分數來實現的

  2. 其次,為了進一步提高訓練效率,減少了在check point的後向傳遞中重新計算的啟用量,在實現上,透過手動實現trasnformer層的後向函式來進行操作

    為了充分受益於這種最佳化,還透過如Korthikanti等人(2022)中採用的方法,進行使用模型和序列並行來減少模型的記憶體使用

  3. 最後,該工作還儘可能地重疊啟用的計算和GPU之間在網路上的通訊
    最終的最佳化效能效果為:當訓練一個65B引數的模型時,程式碼在2048A100的GPU上處理大約380個token/秒/GPU,並耗費80GB的記憶體,這意味著對包含1.4Ttoken的資料集進行訓練大約花費了21天

LLaMA釋出不久後,一些研究者基於它做了不少工作

  • 一開始最小引數7B的模型也需要近30GB的GPU才能執行,但透過bitsandbytes進行浮點最佳化,能夠讓模型在單個NVIDIA RTX 3060(視訊記憶體一般12G)上執行

  • 之後,GitHub 上的一名研究人員甚至能夠在Ryzen 7900X CPU上執行LLM的7B 版本,每秒能推斷出幾個單詞

  • 再之後,有研究者推出了llama.cpp,無需 GPU,就能執行 LLaMA

    llama.cpp 專案實現了在MacBook上執行 LLaMA,還有開發者成功的在 4GB RAM 的樹莓派上執行了 LLaMA 7B

二、LLama2

1. 簡介

23年7月份,Meta釋出LLAMA 2

LLAMA 2 的系列模型有 7B、13B、34B、70B

2. 模型架構

  1. 採用了 Llama 1 的大部分預訓練設定和模型架構,比如使用標準Transformer 架構,使用 RMSNorm 應用Pre-Norm、使用 SwiGLU 啟用函式和旋轉位置嵌入RoPE
  2. 繼續沿用Llama1 所用的透過SentencePiece實現的BPE,且整個詞表的大小依然為 32K
  3. 訓練資料規模是2T個token(即2萬億個token),相比1代的1.4T多出了40%
  4. 上下文長度達到了4096,相比1代的2048直接翻了一倍
  5. 34B以及70B的Llama2應用了分組查詢注意力GQA

2.1 LLaMA 2-Chat:三個版本——7B 13B 70B

  1. 先是監督微調LLaMA2得到SFT版本 (接受了成千上萬個人類標註資料的訓練,本質是問題-答案對 )

  2. 然後使用人類反饋強化學習(RLHF)進行迭代最佳化

    先訓練一個獎勵模型

    然後在獎勵模型/優勢函式的指引下,透過拒絕抽樣(rejection sampling)和近端策略最佳化(PPO)的方法迭代模型的生成策略

    拒絕抽樣是一種從目標分佈中生成樣本的方法

    定義目標分佈和提議分佈

    1.) 目標分佈 p(x):我們希望從中抽樣的複雜分佈。

    2.) 提議分佈 q(x):一個易於從中抽樣的分佈,通常比目標分佈更簡單。

    抽樣和權重計算

    1.) 從提議分佈 q(x)中生成一個樣本 x。

    2.) 計算樣本的權重 \(w(x)=\frac{p(x)}{M q(x)}\),其中 M 是一個大於等於 \(\frac{p(x)}{q(x)}\) 的常數

    接受或拒絕樣本

    1.)生成一個均勻分佈的隨機數 u 在 [0,1] 之間。

    2.)如果 u≤w(x),則接受樣本 x;否則,拒絕樣本並重覆上述過程。

    拒絕抽樣在LLaMA 2-Chat中的應用

    拒絕抽樣:在模型生成答案的過程中,透過獎勵模型評估生成的答案質量。如果答案質量不符合預期,則拒絕該答案並嘗試生成新的答案。這類似於從提議分佈(初始生成的答案)中抽樣,並根據目標分佈(高質量答案的分佈)進行篩選。

LLAMA 2 的效能表現更加接近 GPT-3.5

2.2 分組查詢注意力——Grouped-Query Attention

自迴歸解碼的標準做法是快取序列中先前標記的鍵 (K) 和值 (V) 對,從而加快注意力計算速度
然而,隨著上下文視窗或批次大小的增加,多頭注意力 (MHA)模型中與 KV 快取大小相關的記憶體成本顯著增長

對於較大的模型,KV 快取大小成為瓶頸,鍵和值投影可以在多個頭之間共享,而不會大幅降低效能,可以使用

image-20240719213129799

經實驗論證,GQA 變體在大多數評估任務上的表現與 MHA 基線相當,並且平均優於 MQA 變體

3. Llama 2-Chat中的RLHF

3.1 監督微調(SFT)

  1. 首先重點收集了幾千個高質量 SFT 資料示例,驗證發現效果勝過百萬低質量的資料
  2. 之後發現幾萬次的SFT標註就足以獲得高質量的結果,最終總共收集了27540條用於SFT的標註資料

對模型進行了 2 次微調

???

微調過程中的引數設定

use a cosine learning rate schedule with an initiallearning rate of 2 ×10−5

a weight decay of 0.1,

a batch size of 64,

a sequence length of 4096 token

3.2 訓練兩個獎勵模型:一個偏實用 一個偏安全

Meta 長期以來收集到的獎勵建模資料的統計結果

基於人類應用指定準則的二元比較的大型資料集,也就是獎勵建模資料,近三百萬。

image-20240719213934529

關於獎勵資料:

1.) 與現有的開源資料集相比,Llama2 的偏好資料具有更多的對話回合,平均長度也更長

2.) 獎勵模型將模型響應及其相應的提示(包括前一輪的上下文)作為輸入,並輸出一個標量分數來表示模型生成的質量(例如有用性和安全性)

為了兼顧和平衡模型的實用性和安全性,LLaMA 2團隊訓練了兩個獨立的獎勵模型

1.) 一個針對實用性(稱為實用性RM)進行了最佳化,在內部所有偏實用的獎勵資料集上進行訓練,並結合從內部偏安全的獎勵資料集和開源安全性資料集中統一取樣的同等部分剩餘資料

2.) 另一個針對安全性(安全性RM)進行了最佳化,在內部所有偏安全的獎勵資料和人類無害資料上進行訓練,並以90/10的比例混合內部偏實用的獎勵資料和開源實用性資料

用於下一個token預測的分類頭被替換為用於輸出標量獎勵的迴歸頭

為了使模型行為與人類偏好相一致,Meta 收集了代表了人類偏好經驗取樣的資料,透過針對同一個prompt模型給出的兩個不同的response,人類標註者選擇他們更喜歡的模型輸出。這種人類偏好被用於訓練獎勵模型

\(\mathcal{L}_{\text {ranking }}=-\log \left(\sigma\left(r_{\theta}\left(x, y_{c}\right)-r_{\theta}\left(x, y_{r}\right)\right)\right)\)

為了讓模型可以更好的體會到不同response質量之間的差異,作者團隊將偏好評級被分為4層評級,且考慮到根據這些評級資訊使得獎勵模型對有更多差異的生成,有著不同分數且這些分數上彼此之間的差距儘可能拉開是有用的,為此,在損失中進一步新增一個邊際成分\(m(r)\)

\(\mathcal{L}_{\text {ranking }}=-\log \left(\sigma\left(r_{\theta}\left(x, y_{c}\right)-r_{\theta}\left(x, y_{r}\right)-m(r)\right)\right)\)

邊際 \(m(r)\) 是偏好評級的離散函式

image-20240719215256997

這4個等級是需要有一定的間隔的,彼此之間不能模稜兩可,這個間隔大小是個超引數,可以人為設定,比如小點的間隔1/3或大點的間隔1

2.6 具體的策略迭代:PPO與拒絕取樣

使用兩種主要演算法對 RLHF 進行了微調:

  1. 近端策略最佳化(PPO)

  2. 拒絕取樣(Rejection Sampling)

    模型生成多個回覆後,選擇最佳的回覆作為模型的輸出,過程中,如果生成的回覆不符合預期,就會被拒絕,直到找到最佳回覆從而幫助提高模型的生成質量,使其更符合人類的期望

三、LLama3

1. Llama 3 簡介

Llama 3有兩個版本:8B 和 70B

為了更好的評估llama3的效能,Meta開發了一套新的高質量人類評估集。該評估集包含 1,800 個prompt,涵蓋 12 類任務

2. 模型架構

和Llama 2一樣,Llama 3 繼續採用相對標準的decoder-only transformer架構,但做了如下幾個關鍵的改進

2.1 擴充詞表

Llama 3 使用具有 128K tokens的tokenizer

一方面,分詞器由 SentencePiece 換為了 Tiktoken,與 GPT4 保持一致,可以更有效地對語言進行編碼

另一方面,Token詞表從LLAMA 2的32K擴充到了128K

基準測試顯示,Tiktoken提高了token效率,與 Llama 2 相比,生成的token最多減少了 15% (正由於llama3具有更大的詞表,比llama2的tokenizer具有更大的文字壓縮率)

2.2 採用分組查詢注意力 GQA

為了提高推理效率,Llama 3在 8B 和 70B 都採用了分組查詢注意力(GQA),根據相關實驗可以觀察到,儘管與 Llama 2 7B 相比,模型的引數多了 1B,但改進的分詞器效率和 GQA 有助於保持與 Llama 2 7B 相同的推理效率

image-20240718145201340

多頭注意力 Multi-head: 每個query對應一個key, 對應一個value

分組查詢注意力 Grouped-query: 對query進行分組, 每組對應一個key, 對應一個value

多查詢注意力 Multi-query: 所有的query共用一個key, 共用一個value

值得指出的是,上一個版本 llama 2 的34B和70B才用到了GQA

image-20240718145609197

2.3 上下文擴充套件至8K

在 8192 個token的序列上訓練模型,且透過掩碼操作以確保自注意力不會跨越文件邊界

相比llama 2是一個進步,畢竟llama 2的上下文長度還只有4K

所以如果在平均長度超過4K的任務中使用llama2進行微調, 不得已必須用上longlora/longqlora這類擴充套件長度的技術.

3. 訓練資料

  1. Llama 3 經過超過 15T token 的預訓練 ( 比 Llama 2 使用的資料集大七倍,並且包含四倍多的程式碼,要知道,llama 2的訓練資料才2T個token,即2萬億個token),這些資料全部從公開來源收集

  2. Llama 3 預訓練資料集的中,其中有超過5%的部分由涵蓋 30 多種語言的高質量非英語資料組成。當然,大機率上,這些語言的效能水平不會與英語相同 ( 原因在於其只佔5% )

  3. 為了確保 Llama 3 接受最高質量資料的訓練,他們還開發了一系列資料過濾pipeline。這些pipeline包括使用啟發式過濾器NSFW 過濾器語義重複資料刪除方法文字分類器來預測資料質量

    使用 Llama 2 作為文字質量分類器 為 Llama 3 生成訓練資料

  4. 還進行了廣泛的實驗,以評估在最終預訓練資料集中混合不同來源的資料的最佳方法。這些實驗使能夠選擇一個資料組合,確保 Llama 3 在各種用例(包括瑣事問題、STEM、編碼、歷史知識等)中表現良好

4. 提高預訓練效率

為了有效利用 Llama 3 模型中的預訓練資料,投入了大量精力來擴大預訓練規模。具體來說

  1. 為下游基準評估制定了一系列詳細的縮放法則。這些縮放法則使我們能夠選擇最佳的資料組合,且使我們能夠在實際訓練模型之前預測最大模型在關鍵任務上的效能

    比如在 Llama 3 的開發過程中,對縮放行為進行了一些新的觀察。

    例如,雖然 8B 引數模型的 Chinchilla 最佳訓練計算量對應於約 200B 個token,但發現即使在模型建立之後,模型效能仍在繼續提高接受了兩個數量級以上的資料訓練

    在對多達 15T token進行訓練後,8B 和 70B 引數模型都繼續以對數線性方式改進。較大的模型可以用較少的訓練來匹配這些較小模型的效能,但較小的模型通常是首選,因為它們在推理過程中效率更高

  2. 為了訓練Llama 3的400B的版本,Meta結合了三種型別的並行化:資料並行化、模型並行化和管道並行化

    當同時在 16K GPU 上進行訓練時,可實現每個 GPU 超過 400 TFLOPS 的計算利用率,當然,最終在兩個定製的24K GPU 叢集上進行了訓練

    16K GPU??

    且:

    1.) 為了最大限度地延長 GPU 的正常執行時間,開發了一種先進的新訓練堆疊,可以自動執行錯誤檢測、處理和維護。還極大地改進了硬體可靠性和靜默資料損壞檢測機制

    靜默資料損壞??

    2.) 並且開發了新的可擴充套件儲存系統,以減少檢查點和回滾的開銷。這些改進使總體有效訓練時間超過 95%

    檢查點的開銷??

    回滾的開銷??

綜合起來,這些改進使 Llama 3 的訓練效率比 Llama 2 提高了約三倍

5. 後訓練策略

對指令調整方法進行了創新, 後訓練方法是:監督微調SFT拒絕取樣近端策略最佳化PPO直接策略最佳化DPO的組合

  1. SFT 中使用的prompt質量,以及 PPO 和 DPO 中使用的偏好排名對對齊模型的效能有著巨大的影響

    最終,在模型質量方面的一些最大改進來自於仔細整理這些資料並對人類標註者提供的標註或註釋進行多輪質量保證

  2. 透過 PPO 和 DPO 從偏好排名中學習也極大地提高了 Llama 3 在推理和編碼任務上的效能。

    如果你向模型提出一個它難以回答的推理問題,該模型有時會產生正確的推理軌跡:模型知道如何產生正確的答案,但不知道如何選擇它,但對“偏好排名的訓練”使模型能夠學習如何選擇它

6. Llama 3 的下一步

llama 3中最大的模型有超過 400B 個引數,不過這個模型仍在訓練中

四、擴充套件Llama3至100萬上下文

1. Llama-3-8B 的上下文長度擴充套件到16K

LLaMA 3 8B (base, not instruct)

使用資料集 LongAlpaca https://huggingface.co/datasets/Yukang/LongAlpaca-16k-length

rope_theta ( base ) was set to 1000000.0 ( 一百萬 )

方法: 針對位置編碼的base引數(rope_theta)擴大

實現步驟:

  1. 首先微調得到一個加長版的模型

    llama3的rope_theta設定為50萬, 這裡擴大到了100萬

    rope_theta引數其實是RoPE中的base, 也就是之前大多模型設定為1萬的引數, 並不是旋轉角度 \(\theta\)

    RoPE的構造基礎是Sinusoidal位置編碼, 可以改寫為下面的公式

    \(\left[\cos \left(\frac{n}{\beta^{0}}\right), \sin \left(\frac{n}{\beta^{0}}\right), \cos \left(\frac{n}{\beta^{1}}\right), \sin \left(\frac{n}{\beta^{1}}\right), \cdots, \cos \left(\frac{n}{\beta^{d / 2-1}}\right), \sin \left(\frac{n}{\beta^{d / 2-1}}\right)\right]\)

    其中 \(\beta = 10000^{\frac{2}{d}}\), 這個\(10000\)即為base

    對base做放大就是ntk-aware插值的操作

    這裡將llama3的rope_theta從50萬放大到100萬, 就是 \(\alpha = 2\) 的ntk-aware插值.

  2. 有了擴充套件好上下文的微調模型之後,使用開源工具Mergekit比較微調模型和基礎模型,提取引數的差異成為LoRA

  3. 同樣使用Mergekit,把提取好的LoRA合併到其他同架構模型中

2. LLama3長度擴充套件到100萬

仍採用NTK-aware插值方法,將rope_theta (base) 繼續放大, 可以使其長度達到100萬

具體方法是:

  1. llama3初始的base值為50萬, 上下文長度為8k

    base值擴大2倍, 上下文長度為16k

    上下文長度要從8k擴充套件到100萬, base值需要擴大125倍, 即 (\(50萬 \times 125 = 6250萬\))

    長度擴充套件多少倍則對應的這個rope_theta擴大多少倍

  2. 漸進式訓練: 使用UC伯克利Pieter Abbeel團隊提出的Blockwise RingAttention方法擴充套件模型的上下文長度

    且團隊透過自定義網路拓撲在Ring Attention之上分層並行化,更好地利用大型GPU叢集來應對裝置之間傳遞許多KV blocks帶來的網路瓶頸,最終使模型的訓練速度提高了33倍

相關文章