在 Windows 上运行 Llama3 源代码

349 阅读2分钟

该模型需要运行在 GPU 上,需要根据显卡安装对应的 CUDA 驱动。安装 CUDA 的步骤可以参考其他 jym 写的文章,此处提供个参考链接 windows操作系统安装配置CUDA12.3和cuDNN

下载代码

Github Llama3 官方仓库

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

下载模型

魔搭社区 Llama-3-8B-Instruct 下载资源

git lfs install 
git clone https://www.modelscope.cn/wdndev/Meta-Llama-3-8B-Instruct-torch.git

备注:上述提供的魔搭社区 Llama-3-8B-Instruct 资源可以无需注册直接下载, 也可以点击 Meta Llama websiteHugging Face 下载模型。

安装依赖

在项目根目录执行以下命令安装依赖。(建议使用 venv 或 conda 避免依赖冲突)

pip install -e .

备注:

  1. 安装完依赖后根据 torch 版本,通过 PyTorch官网 找到对应 cuda 版本的 torch 并重新安装。
  2. 运行代码时可能会报找不到指定的模块, 通过命令 pip install XXXX 安装缺失的依赖即可。

运行代码

项目中新建文件并复制以下代码,并修改变量 ckpt_dir 和 tokenizer_path 成下载完成的模型路径。

import torch
import os
from typing import List, Optional
from llama import Dialog, Llama

def init_torch():
    os.environ['MASTER_ADDR'] = 'localhost'
    # 自定义的端口号,可修改成不冲突的端口号都行
    os.environ['MASTER_PORT'] = '19999'
    #os.environ["CUDA_VISIBLE_DEVICES"] = '0'
    os.environ["RANK"] = '0'
    os.environ["WORLD_SIZE"] = '1'
    
    if not torch.distributed.is_initialized():
        torch.distributed.init_process_group(backend='gloo')

def build_llama(
    ckpt_dir: str,
    tokenizer_path: str,
    max_seq_len: int = 512,
    max_batch_size: int = 4) -> Llama:
    
    generator = Llama.build(
        ckpt_dir=ckpt_dir,
        tokenizer_path=tokenizer_path,
        max_seq_len=max_seq_len,
        max_batch_size=max_batch_size,
        model_parallel_size=1
    )

    return generator

def chat_completion(
    generator: Llama,     
    dialogs: List[Dialog],
    temperature: float = 0.6,
    top_p: float = 0.9,
    max_gen_len: Optional[int] = None):
    
    results = generator.chat_completion(
        dialogs,
        max_gen_len=max_gen_len,
        temperature=temperature,
        top_p=top_p,
    )
    
    return results

if __name__ == "__main__":
    ckpt_dir = "C:\\model-repos\\Meta-Llama-3-8B-Instruct-torch"
    tokenizer_path = "C:\\model-repos\\Meta-Llama-3-8B-Instruct-torch\\tokenizer.model"
    dialogs: List[Dialog] = [
        [
            {"role": "system", "content": "Always answer with Haiku"},
            {"role": "user", "content": "I am going to Paris, what should I see?"},
        ]
    ]
    
    # 初始化 PyTorch Distributed,以支持 Windows 上运行 Llama3 代码
    init_torch()
    # 创建 Llama 模型实例
    generator = build_llama(ckpt_dir=ckpt_dir, tokenizer_path=tokenizer_path)
    # 调用方法进行对话补全
    results = chat_completion(generator=generator, dialogs=dialogs)
    
    for dialog, result in zip(dialogs, results):
        for msg in dialog:
            print(f"{msg['role'].capitalize()}: {msg['content']}\n")
        print(
            f"> {result['generation']['role'].capitalize()}: {result['generation']['content']}"
        )
        print("\n==================================\n")

运行后结果大概如下:

image.png

参考资料