2023年的深度學(xué)習(xí)入門指南(19) - LLaMA 2源碼解析

上一節(jié)我們學(xué)習(xí)了LLaMA 2的補全和聊天兩種API的使用方法。本節(jié)我們來看看LLaMA 2的源碼。

補全函數(shù)text_completion源碼解析

上一節(jié)我們講了LLaMA 2的編程方法。我們來復(fù)習(xí)一下：

    generator = Llama.build(
        ckpt_dir=ckpt_dir,
        tokenizer_path=tokenizer_path,
        max_seq_len=max_seq_len,
        max_batch_size=max_batch_size,
    )

    prompts = [
        "上下五千年，英雄萬萬千。黃沙百戰(zhàn)穿金甲，不破樓蘭終不還",
    ]
    results = generator.text_completion(
        prompts,
        max_gen_len=max_gen_len,
        temperature=temperature,
        top_p=top_p,
    )

我們先來看看text_completion函數(shù)的參數(shù)是什么意思，該函數(shù)的原型為：

def text_completion(
self,
prompts: List[str],
temperature: float = 0.6,
top_p: float = 0.9,
max_gen_len: Optional[int] = None,
logprobs: bool = False,
echo: bool = False,
) -> List[CompletionPrediction]:

我們來看下這些參數(shù)的含義：

• prompts：這是一個字符串列表，每個字符串都是一個用于生成文本的提示。
• temperature（默認(rèn)值為0.6）：這是一個控制生成文本隨機性的參數(shù)。溫度值越高，生成的文本就越隨機；溫度值越低，生成的文本就越傾向于最可能的輸出。
• top_p（默認(rèn)值為0.9）：這是一個控制生成文本多樣性的參數(shù)，它設(shè)定了從最高概率的詞開始，累計到總概率超過top_p的詞為止，然后從這些詞中隨機選擇一個詞作為生成的詞。這種方法也被稱為nucleus sampling或top-p sampling。
• max_gen_len：可選參數(shù)，表示生成的文本的最大長度。如果未指定，那么將使用模型參數(shù)中的最大序列長度減1。
• logprobs（默認(rèn)值為False）：如果為True，那么在返回的結(jié)果中會包含生成的每個詞的對數(shù)概率。
• echo（默認(rèn)值為False）：這是一個控制是否在生成的文本中包含輸入提示的參數(shù)。

參數(shù)明白了之后我們看text_completion完整實現(xiàn)：

    def text_completion(
        self,
        prompts: List[str],
        temperature: float = 0.6,
        top_p: float = 0.9,
        max_gen_len: Optional[int] = None,
        logprobs: bool = False,
        echo: bool = False,
    ) -> List[CompletionPrediction]:
        if max_gen_len is None:
            max_gen_len = self.model.params.max_seq_len - 1
        prompt_tokens = [self.tokenizer.encode(x, bos=True, eos=False) for x in prompts]
        generation_tokens, generation_logprobs = self.generate(
            prompt_tokens=prompt_tokens,
            max_gen_len=max_gen_len,
            temperature=temperature,
            top_p=top_p,
            logprobs=logprobs,
            echo=echo,
        )
        if logprobs:
            return [
                {
                    "generation": self.tokenizer.decode(t),
                    "tokens": [self.tokenizer.decode(x) for x in t],
                    "logprobs": logprobs_i,
                }
                for t, logprobs_i in zip(generation_tokens, generation_logprobs)
            ]
        return [{"generation": self.tokenizer.decode(t)} for t in generation_tokens]

總結(jié)起來就三步，這個text_completion其實就是generate的包裝函數(shù)：

• 編碼：調(diào)用tokenizer.encode
• 生成：調(diào)用generate
• 解碼：調(diào)用tokenizer.decode

分詞

import os
from logging import getLogger
from typing import List

from sentencepiece import SentencePieceProcessor


logger = getLogger()


class Tokenizer:
    def __init__(self, model_path: str):
        # reload tokenizer
        assert os.path.isfile(model_path), model_path
        self.sp_model = SentencePieceProcessor(model_file=model_path)
        logger.info(f"Reloaded SentencePiece model from {model_path}")

        # BOS / EOS token IDs
        self.n_words: int = self.sp_model.vocab_size()
        self.bos_id: int = self.sp_model.bos_id()
        self.eos_id: int = self.sp_model.eos_id()
        self.pad_id: int = self.sp_model.pad_id()
        logger.info(
            f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"
        )
        assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()

    def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
        assert type(s) is str
        t = self.sp_model.encode(s)
        if bos:
            t = [self.bos_id] + t
        if eos:
            t = t + [self.eos_id]
        return t

    def decode(self, t: List[int]) -> str:
        return self.sp_model.decode(t)

首先是用到了分詞組件SentencePieceProcessor。SentencePieceProcessor是SentencePiece庫中的一個組件,它實現(xiàn)了子詞(subword)tokenize和detokenize的功能。

其主要作用包括:

• 將文本tokenize成子詞(subword)。SentencePiece 使用的數(shù)據(jù)驅(qū)動方法,可以學(xué)習(xí)文本的詞匯表并將文本tokenize成子詞單元。
• 將子詞detokenize合并成原始文本?？梢詫okenize后的子詞序列重新合并為原始文本。
• 提供vocab管理?？梢垣@得tokenize的子詞詞匯表等信息。
• 支持多種語言文本的tokenize和detokenize。
• 提供高效的實現(xiàn)。底層使用C++實現(xiàn),可以快速處理大規(guī)模文本。
• 提供多種模型選擇,如BPE、unigram等。
• 支持自定義訓(xùn)練子詞化模型。

好，下面我們回到這段代碼本身。這段代碼實現(xiàn)了一個基于SentencePiece的Tokenizer類,可以進行文本的tokenize和detokenize。

主要邏輯:

• 在初始化時加載SentencePiece模型文件model_path。
• 獲取模型的詞匯表大小n_words,以及特殊token的id(bos_id,eos_id,pad_id)。
• encode方法可以將字符串文本s tokenize成id列表。可以選擇在開始加入bos_id,結(jié)尾加入eos_id。
• decode方法可以將id列表解碼還原為字符串文本。

這樣就構(gòu)建了一個封裝SentencePiece tokenize/detokenize的Tokenizer類?？梢约虞d自定義的SentencePiece模型,然后就可以方便地對文本進行子詞化處理。

這種方式可以重復(fù)使用已訓(xùn)練好的SentencePiece模型,為下游NLP任務(wù)提供可靠的tokenize和detokenize功能。

最后我們再講一講幾個特殊的符號bos_id、eos_id和pad_id：

• bos_id: 開始符(Beginning of Sentence)的id。用于表示一個序列的開始。
• eos_id: 結(jié)束符(End of Sentence)的id。用于表示一個序列的結(jié)束。
• pad_id: 填充符(Padding)的id。當(dāng)需要將多個序列長度對齊時,可以使用pad_id在較短序列后面填充。

聊天函數(shù)chat_completion

在進入generate函數(shù)之前，我們再看看chat_completion是如何實現(xiàn)的。

    def chat_completion(
        self,
        dialogs: List[Dialog],
        temperature: float = 0.6,
        top_p: float = 0.9,
        max_gen_len: Optional[int] = None,
        logprobs: bool = False,
    ) -> List[ChatPrediction]:
        if max_gen_len is None:
            max_gen_len = self.model.params.max_seq_len - 1
        prompt_tokens = []
        for dialog in dialogs:
            if dialog[0]["role"] != "system":
                dialog = [
                    {
                        "role": "system",
                        "content": DEFAULT_SYSTEM_PROMPT,
                    }
                ] + dialog
            dialog = [
                {
                    "role": dialog[1]["role"],
                    "content": B_SYS
                    + dialog[0]["content"]
                    + E_SYS
                    + dialog[1]["content"],
                }
            ] + dialog[2:]
            assert all([msg["role"] == "user" for msg in dialog[::2]]) and all(
                [msg["role"] == "assistant" for msg in dialog[1::2]]
            ), (
                "model only supports 'system', 'user' and 'assistant' roles, "
                "starting with 'system', then 'user' and alternating (u/a/u/a/u...)"
            )
            dialog_tokens: List[int] = sum(
                [
                    self.tokenizer.encode(
                        f"{B_INST} {(prompt['content']).strip()} {E_INST} {(answer['content']).strip()} ",
                        bos=True,
                        eos=True,
                    )
                    for prompt, answer in zip(
                        dialog[::2],
                        dialog[1::2],
                    )
                ],
                [],
            )
            assert (
                dialog[-1]["role"] == "user"
            ), f"Last message must be from user, got {dialog[-1]['role']}"
            dialog_tokens += self.tokenizer.encode(
                f"{B_INST} {(dialog[-1]['content']).strip()} {E_INST}",
                bos=True,
                eos=False,
            )
            prompt_tokens.append(dialog_tokens)

        generation_tokens, generation_logprobs = self.generate(
            prompt_tokens=prompt_tokens,
            max_gen_len=max_gen_len,
            temperature=temperature,
            top_p=top_p,
            logprobs=logprobs,
        )
        if logprobs:
            return [
                {
                    "generation": {
                        "role": "assistant",
                        "content": self.tokenizer.decode(t),
                    },
                    "tokens": [self.tokenizer.decode(x) for x in t],
                    "logprobs": logprobs_i,
                }
                for t, logprobs_i in zip(generation_tokens, generation_logprobs)
            ]
        return [
            {"generation": {"role": "assistant", "content": self.tokenizer.decode(t)}}
            for t in generation_tokens
        ]

我們先看一下參數(shù)：

• dialogs：一個對話列表，其中每個對話都是一個字典列表，表示一段對話。
• temperature：一個浮點數(shù)，表示生成文本時使用的溫度。默認(rèn)值為 0.6。
• top_p：一個浮點數(shù)，表示生成文本時使用的 top-p 采樣。默認(rèn)值為 0.9。
• max_gen_len：一個可選的整數(shù)，表示生成文本的最大長度。如果未指定，則使用模型參數(shù)中的最大序列長度減一。
• logprobs：一個布爾值，表示是否返回生成文本的對數(shù)概率。默認(rèn)值為 False。
  函數(shù)返回一個 ChatPrediction 列表，其中每個元素都是一個字典，包含生成的回復(fù)和相關(guān)信息。

函數(shù)首先檢查 max_gen_len 是否為 None，如果是，則將其設(shè)置為模型參數(shù)中的最大序列長度減一。然后，對于每個對話，函數(shù)執(zhí)行以下操作：

• 如果第一條消息的角色不是 “system”，則在對話的開頭添加一條默認(rèn)的系統(tǒng)提示。
• 將第一條和第二條消息合并為一條消息，并更新對話。
• 檢查對話中消息的角色是否符合要求（即以 “system” 開始，然后交替出現(xiàn) “user” 和 “assistant”）。
• 對于每一組相鄰的提示和回答（即每兩條消息），使用 tokenizer 對其進行編碼，并將編碼后的 token 連接起來。
• 檢查最后一條消息是否來自用戶。
• 對最后一條消息進行編碼，并將編碼后的 token 添加到 token 列表中。

接下來，函數(shù)調(diào)用 generate 方法生成回復(fù)，并根據(jù) logprobs 參數(shù)的值返回相應(yīng)的結(jié)果。如果 logprobs 為 True，則返回包含生成回復(fù)、token 和對數(shù)概率的字典列表；否則，返回僅包含生成回復(fù)的字典列表。這些生成回復(fù)都具有 “assistant” 角色，并使用 tokenizer 進行解碼。

總體來說，只是增加了對于對話角色的業(yè)務(wù)邏輯處理，核心還是調(diào)用generate函數(shù)。

溫度與top p采樣

在進入講解generate函數(shù)之前，我們先講一個小知識點，就是溫度temperature的作用。我們看下面的代碼：

            if temperature > 0:
                probs = torch.softmax(logits[:, -1] / temperature, dim=-1)
                next_token = sample_top_p(probs, top_p)
            else:
                next_token = torch.argmax(logits[:, -1], dim=-1)

temperature 是一個超參數(shù)，用于控制生成文本的多樣性。當(dāng) temperature 較高時，概率分布更加平坦，因此采樣出的標(biāo)記更具多樣性。當(dāng) temperature 較低時，概率分布更加尖銳，因此采樣出的標(biāo)記更傾向于概率最大的那個。當(dāng) temperature 等于 0 時，直接選擇概率最大的標(biāo)記。

那么，sample_top_p是如何實現(xiàn)的呢？我把解說寫在代碼注釋里面了：

def sample_top_p(probs, p):
    # 這行代碼將輸入的概率 probs 按照降序排序。probs_sort 是排序后的概率，probs_idx 是對應(yīng)的索引。
    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
    # 這行代碼計算 probs_sort 的累積和。累積和是從第一個元素開始，依次將序列中的每個元素與前面所有元素的和相加得到的。
    probs_sum = torch.cumsum(probs_sort, dim=-1)
    # 這行代碼生成一個布爾掩碼，用于指示哪些累積和減去當(dāng)前概率的值大于 p。這用于確定哪些概率應(yīng)該被設(shè)為0，以保證被抽樣的概率和不超過 p。
    mask = probs_sum - probs_sort > p
    # 這行代碼使用上述生成的掩碼，將那些使累積和減去當(dāng)前概率的值大于 p 的 probs_sort 中的元素設(shè)為0。
    probs_sort[mask] = 0.0
    # 這行代碼將 probs_sort 中的每個元素除以它們的和，以便重新歸一化概率分布。
    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
    # 這行代碼從歸一化的 probs_sort 中抽取一個樣本。torch.multinomial 是PyTorch中的多項式分布抽樣函數(shù)，它根據(jù)每個元素的權(quán)重抽取樣本。
    next_token = torch.multinomial(probs_sort, num_samples=1)
    # 這行代碼使用 torch.gather 函數(shù)從 probs_idx 中收集對應(yīng) next_token 的索引，這樣就能得到原始概率 probs 中對應(yīng)的索引。
    next_token = torch.gather(probs_idx, -1, next_token)
    return next_token

總的來說，sample_top_p保留了按概率高低排序的大致分布,但過濾了長尾部分的低概率噪聲。然后從重歸一化的分布中采樣,既保證了質(zhì)量,又增加了適當(dāng)?shù)碾S機性。

generate函數(shù)

好，我們終于開始探索最核心的生成函數(shù)上了：

    @torch.inference_mode()
    def generate(
        self,
        prompt_tokens: List[List[int]],
        max_gen_len: int,
        temperature: float = 0.6,
        top_p: float = 0.9,
        logprobs: bool = False,
        echo: bool = False,
    ) -> Tuple[List[List[int]], Optional[List[List[float]]]]:

首先是這個函數(shù)的參數(shù)，其實我們已經(jīng)比較熟悉了。包括輸入的提示 tokens（prompt_tokens），最大生成長度（max_gen_len），溫度參數(shù)（temperature，影響生成文本的隨機性）， top_p（用于決定采樣過程中保留的 token 集合的概率閾值，也被稱為 “nucleus sampling”），是否返回每個 token 的對數(shù)概率（logprobs），以及是否將輸入的提示返回（echo）。

        params = self.model.params
        bsz = len(prompt_tokens)
        assert bsz <= params.max_batch_size, (bsz, params.max_batch_size)

        min_prompt_len = min(len(t) for t in prompt_tokens)
        max_prompt_len = max(len(t) for t in prompt_tokens)
        assert max_prompt_len <= params.max_seq_len
        total_len = min(params.max_seq_len, max_gen_len + max_prompt_len)

        pad_id = self.tokenizer.pad_id
        tokens = torch.full((bsz, total_len), pad_id, dtype=torch.long, device="cuda")
        for k, t in enumerate(prompt_tokens):
            tokens[k, : len(t)] = torch.tensor(t, dtype=torch.long, device="cuda")
        if logprobs:
            token_logprobs = torch.zeros_like(tokens, dtype=torch.float)

        prev_pos = 0
        eos_reached = torch.tensor([False] * bsz, device="cuda")
        input_text_mask = tokens != pad_id

接著，根據(jù)提供的 prompt_tokens 初始化一個 tokens 張量，長度為 total_len，并填充模型的 pad_id。然后，將 prompt_tokens 的內(nèi)容復(fù)制到 tokens 張量的對應(yīng)位置。

        for cur_pos in range(min_prompt_len, total_len):
            logits = self.model.forward(tokens[:, prev_pos:cur_pos], prev_pos)
            if logprobs:
                token_logprobs[:, prev_pos + 1 : cur_pos + 1] = -F.cross_entropy(
                    input=logits.transpose(1, 2),
                    target=tokens[:, prev_pos + 1 : cur_pos + 1],
                    reduction="none",
                    ignore_index=pad_id,
                )
            if temperature > 0:
                probs = torch.softmax(logits[:, -1] / temperature, dim=-1)
                next_token = sample_top_p(probs, top_p)
            else:
                next_token = torch.argmax(logits[:, -1], dim=-1)

            next_token = next_token.reshape(-1)
            # only replace token if prompt has already been generated
            next_token = torch.where(
                input_text_mask[:, cur_pos], tokens[:, cur_pos], next_token
            )
            tokens[:, cur_pos] = next_token
            eos_reached |= (~input_text_mask[:, cur_pos]) & (
                next_token == self.tokenizer.eos_id
            )
            prev_pos = cur_pos
            if all(eos_reached):
                break

然后，對于 tokens 張量中的每一個位置，計算下一個 token 的 logits，并基于這些 logits 生成下一個 token。如果 logprobs 參數(shù)為真，則計算每個 token 的對數(shù)概率。如果溫度大于 0，則使用 softmax 函數(shù)和溫度參數(shù)對 logits 進行縮放，然后使用 top-p 采樣生成下一個 token。否則，直接選擇 logits 最大的 token。新生成的 token 會替換 tokens 張量中的對應(yīng)位置。

如果生成的 token 是結(jié)束標(biāo)記（eos_id），則更新 eos_reached 標(biāo)記。如果所有的序列都已經(jīng)生成了結(jié)束標(biāo)記，則停止生成。

        if logprobs:
            token_logprobs = token_logprobs.tolist()
        out_tokens, out_logprobs = [], []
        for i, toks in enumerate(tokens.tolist()):
            # cut to max gen len
            start = 0 if echo else len(prompt_tokens[i])
            toks = toks[start : len(prompt_tokens[i]) + max_gen_len]
            probs = None
            if logprobs:
                probs = token_logprobs[i][start : len(prompt_tokens[i]) + max_gen_len]
            # cut to eos tok if any
            if self.tokenizer.eos_id in toks:
                eos_idx = toks.index(self.tokenizer.eos_id)
                toks = toks[:eos_idx]
                probs = probs[:eos_idx] if logprobs else None
            out_tokens.append(toks)
            out_logprobs.append(probs)
        return (out_tokens, out_logprobs if logprobs else None)

最后，如果 logprobs 參數(shù)為真，則將 token_logprobs 轉(zhuǎn)換為列表。然后，對于 tokens 張量中的每一行（即每一個生成的序列），如果 echo 參數(shù)為假，則去掉提示部分。然后，如果存在結(jié)束標(biāo)記，則去掉結(jié)束標(biāo)記之后的部分。最后，返回生成的 tokens 和對數(shù)概率（如果 logprobs 參數(shù)為真）。

這個函數(shù)返回的是一個元組，第一個元素是一個列表，包含每一個生成的 token 序列。第二個元素是一個列表，包含每一個生成的對數(shù)概率序列（如果 logprobs 參數(shù)為真）。

build構(gòu)造函數(shù)

最后我們再說下構(gòu)造Llama的部分：

    @staticmethod
    def build(
        ckpt_dir: str,
        tokenizer_path: str,
        max_seq_len: int,
        max_batch_size: int,
        model_parallel_size: Optional[int] = None,
    ) -> "Llama":
        if not torch.distributed.is_initialized():
            torch.distributed.init_process_group("nccl")
        if not model_parallel_is_initialized():
            if model_parallel_size is None:
                model_parallel_size = int(os.environ.get("WORLD_SIZE", 1))
            initialize_model_parallel(model_parallel_size)

        local_rank = int(os.environ.get("LOCAL_RANK", 0))
        torch.cuda.set_device(local_rank)

        # seed must be the same in all processes
        torch.manual_seed(1)

        if local_rank > 0:
            sys.stdout = open(os.devnull, "w")

        start_time = time.time()
        checkpoints = sorted(Path(ckpt_dir).glob("*.pth"))
        assert len(checkpoints) > 0, f"no checkpoint files found in {ckpt_dir}"
        assert model_parallel_size == len(
            checkpoints
        ), f"Loading a checkpoint for MP={len(checkpoints)} but world size is {model_parallel_size}"
        ckpt_path = checkpoints[get_model_parallel_rank()]
        checkpoint = torch.load(ckpt_path, map_location="cpu")
        with open(Path(ckpt_dir) / "params.json", "r") as f:
            params = json.loads(f.read())

        model_args: ModelArgs = ModelArgs(
            max_seq_len=max_seq_len,
            max_batch_size=max_batch_size,
            **params,
        )
        tokenizer = Tokenizer(model_path=tokenizer_path)
        model_args.vocab_size = tokenizer.n_words
        torch.set_default_tensor_type(torch.cuda.HalfTensor)
        model = Transformer(model_args)
        model.load_state_dict(checkpoint, strict=False)
        print(f"Loaded in {time.time() - start_time:.2f} seconds")

        return Llama(model, tokenizer)

雖然這么一大段，但其實都是一些初始化的工作。

• 分布式設(shè)置：首先，這段代碼檢查是否已經(jīng)初始化了 PyTorch 的分布式環(huán)境，如果沒有則進行初始化。然后，檢查是否已經(jīng)初始化了模型并行環(huán)境，如果沒有，則獲取環(huán)境變量 WORLD_SIZE 的值作為模型并行的大小，并進行初始化。

• 設(shè)備設(shè)置：獲取環(huán)境變量 LOCAL_RANK 的值作為本地排名，并設(shè)置當(dāng)前設(shè)備為該排名對應(yīng)的 GPU。

• 隨機種子設(shè)置：為了確保所有進程生成的隨機數(shù)相同，設(shè)置隨機種子為 1。

• 標(biāo)準(zhǔn)輸出設(shè)置：如果本地排名大于 0，則將標(biāo)準(zhǔn)輸出重定向到空設(shè)備，即不顯示任何輸出。

• 加載模型檢查點：找到檢查點目錄中的所有檢查點文件，并按照文件名排序。然后，根據(jù)模型并行的排名選擇一個檢查點文件，并加載該檢查點。然后，加載模型參數(shù)。

• 構(gòu)建模型和分詞器：使用加載的模型參數(shù)和提供的 max_seq_len 和 max_batch_size 構(gòu)建模型參數(shù)對象。然后，加載分詞器，并設(shè)置模型參數(shù)的詞匯表大小為分詞器的詞匯表大小。然后，設(shè)置默認(rèn)的張量類型為半精度浮點型（以節(jié)省內(nèi)存和計算資源）。然后，構(gòu)建 Transformer 模型，并加載模型檢查點。

• 最后，構(gòu)建一個 Llama 對象，包含加載的模型和分詞器，并返回該對象。

小結(jié)

本節(jié)我們學(xué)習(xí)了LLaMA 2的源碼，包括補全函數(shù)text_completion和聊天函數(shù)chat_completion的實現(xiàn)，以及它們的真正實現(xiàn)generate函數(shù)的原理。我們還學(xué)習(xí)了溫度temperature和top p采樣的原理。

對于沒有搞到深度學(xué)習(xí)生成的同學(xué)，可能有一點難度。

LLaMA的代碼還差一部分就是模型的部分，我們放到下一節(jié)來講。要不然知識點太多大家容易大腦缺氧：）文章來源地址http://www.zghlxwxcb.cn/news/detail-616145.html

到了這里，關(guān)于2023年的深度學(xué)習(xí)入門指南(19) - LLaMA 2源碼解析的文章就介紹完了。如果您還想了解更多內(nèi)容，請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！