阅读llama源码笔记_1

文章目录

总体逻辑

以如下命令为例，梳理一下总体流程

torchrun --nproc_per_node 1 example_text_completion.py \
    --ckpt_dir llama-2-7b-chat/ \
    --tokenizer_path tokenizer.model \
    --max_seq_len 512 --max_batch_size 6

在顶层目录运行使用torchrun运行example_text_completion.py文件，example_text_completion.py中逻辑可分为三大步：

1) 构造模型

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

2) 定义数据

prompts: List[str] = [
        # For these prompts, the expected answer is the natural continuation of the prompt
        "I believe the meaning of life is",
        "Simply put, the theory of relativity states that ",
        """A brief message congratulating the team on the launch:

        Hi everyone,
        
        I just """,
        # Few shot prompt (providing a few examples before asking model to complete more);
        """Translate English to French:
        
        sea otter => loutre de mer
        peppermint => menthe poivrée
        plush girafe => girafe peluche
        cheese =>""",
    ]

3) 处理下游任务

对不同的数据有不同的处理方法，对相同的数据也有不同的任务，text_completion函数处理的下游任务是文本生成

results = generator.text_completion(
        prompts,
        max_gen_len=max_gen_len,
        temperature=temperature,
        top_p=top_p,
    )

接下来对三个步骤进行分析，只列出关键代码

1.构造模型

Llama.build函数是llama/llama/目录下generation.py文件中定义的，build最开始的一些并行参数设置就不分析了，因为我目前也不是很懂，在单GPU上跑可以先不考虑并行的相关操作。

build函数构建了一个tokenizer模型（分词器）和一个Transformer模型，分别对应tokenizer.py文件和model.py文件，其实也很好理解，输入的字符串要先过分词器变为tokenid再输入给Transformer

build函数逻辑上也可以分为2步：

1.1构造分词器tokenizer

tokenizer = Tokenizer(model_path=tokenizer_path)

tokenizer是放在CPU上的

tokenizer.encode()将一个字符串编码成一串tokenId

tokenizer.decode() 一串tokenId解码成字符串

1.2构造Transformer

先调用torch.load将权重参数加载到CPU上，然后在GPU上初始化Transformer模型，最后再调用load_state_dict将CPU上的参数传输到GPU赋值给Transformer模型

checkpoint = torch.load(ckpt_path, map_location="cpu")//将权重参数加载到CPU上
torch.set_default_tensor_type(torch.cuda.HalfTensor)
model = Transformer(model_args)//在GPU上初始化Transformer模型    底层调用cudaMalloc在GPU开辟参数存放空间，并使用FillFunctor将参数都初始化为0
model.load_state_dict(checkpoint, strict=False)//将CPU上的参数复制到GPU赋值给Transformer模型  底层成对调用cudaMemcpyAsync和cudaStreamSynchronize

底层的调用关系我是通过nsys computer进行分析的

2.定义数据

没啥好讲的，声明一个 List[str]对象，当然也可以从网上下载其他数据集，如果真正线上部署的话，那么就是通过网络请求传来的数据了

3.处理下游任务

将数据通过generator.text_completion函数(generation.py中定义的)处理，text_completion函数逻辑如下：

3.1分词（编码）

将字符串编码为tokenid

prompt_tokens = [self.tokenizer.encode(x, bos=True, eos=False) for x in prompts]

3.2 generate

该函数逻辑后面再分析，函数的返回值就是生成的tokenId

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

3.3 解码

即将tokenId解码成字符串，即该下游任务的最终结果

self.tokenizer.decode(t)} for t in generation_tokens