Llama3 本地部署 CPU进行推理

下载模型

提交邮件就能申请，在官方GitHub仓库下载示例代码进行使用。

https://github.com/meta-llama/llama3

GPU推理

Llama3默认是开启组并行推理，所以需要显卡来运行，参考官方示例很容易便完成的运行。

  1. 加载llama模型，
  2. 构造prompt
  3. 输入llama模型得到结果
  4. 处理结果
def main(
    ckpt_dir: str,
    tokenizer_path: str,
    temperature: float = 0.6,
    top_p: float = 0.9,
    max_seq_len: int = 128,
    max_gen_len: int = 64,
    max_batch_size: int = 4,
):
    generator = Llama.build(
        ckpt_dir=ckpt_dir,
        tokenizer_path=tokenizer_path,
        max_seq_len=max_seq_len,
        max_batch_size=max_batch_size,
    )
    
    prompts: List[str] = [
        """Translate English to Chinese:

        cheese =>""",
    ]
    results = generator.text_completion(
        prompts,
        max_gen_len=max_gen_len,
        temperature=temperature,
        top_p=top_p,
    )
    
    for prompt, result in zip(prompts, results):
        print(prompt)
        print(f"> {result['generation']}")
        print("\n==================================\n")

CPU

既然Llama3是默认是开始GPU运行的，那我们能不能在CPU上部署运行，答案是能行的，只能需要改到代码。

  • fairscale里面的parallel全部退化为nn里面提供的
  • 代码里默认cuda，抽离其中的device，放到函数中
  • 退化过程中，某些参数进行删除，依据报错信息来
  • 去掉DDP

完成上述操作之后，就能在CPU上推理Llama3，我测试了下在A100服务器的CPU上推理一次大概1分钟。