运行gpt2预训练库

0x00 TLDR;

1、了解大模型文件加载和执行的一般过程
2、了解模型文件本地存储的其中一种形式。

0x01 环境&准备

  • 系统环境
cat /etc/os-release
# python3 --version
Python 3.10.12

# nvidia-smi
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 575.51.03              Driver Version: 575.51.03      CUDA Version: 12.9     |

# pip3 list | grep torch
torch                    2.7.1
  • 手动下载模型

由于原始模型文件无法直接访问,采用国内镜像方式预先手动下载。

# 安装下载工具
pip3 install huggingface-hub

# 配置加速镜像
echo 'export HF_ENDPOINT="https://hf-mirror.com"' >> ~/.bashrc

# 下载文件至默认缓存路径:  ~/.cache/huggingface/hub
huggingface-cli download "openai-community/gpt2" --include config.json pytorch_model.bin vocab.json merges.txt tokenizer_config.json

# 下载成功后,看到的 hash 目录就是模型路径
ls ~/.cache/huggingface/hub/models--openai-community--gpt2/snapshots

0x02 代码


import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def poc_transformer_model():
    # 这里如果报错的话,使用绝对路径即可
    model_dir = "~/.cache/huggingface/hub/models--openai-community--gpt2/snapshots/{上述的hash目录}"
    tokenizer = AutoTokenizer.from_pretrained(model_dir, local_files_only=True)
    model = AutoModelForCausalLM.from_pretrained(model_dir, local_files_only=True)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = model.to(device)

    prompt = "Why sky is blue?"

    tokenizer.pad_token = tokenizer.eos_token
    # Encode
    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        padding=False,
        truncation=True,
        return_attention_mask=True
    ).to(device)

    # Output
    outputs = model.generate(
        inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        max_length=50,
        num_return_sequences=3,
        do_sample=True,
        top_k=50,
        top_p=0.95,
        temperature=0.8,
        pad_token_id=tokenizer.eos_token_id
    )

    for i, output in enumerate(outputs):
        agc = tokenizer.decode(output, skip_special_tokens=True)
        print(f'生成的内容{i + 1}: {agc}')


if __name__ == '__main__':
    print(torch.__version__)
    print(torch.cuda.is_available())

    poc_transformer_model()

0x03 监控工具

常用GPU使用工具有 gpustat, nvtop

  • 安装 sudo apt install gpustat nvtop
  • 使用
# 启动后自动刷新,可以看历史
nvtop

# 每秒刷新一次
gpustat -i 1
最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
平台声明:文章内容(如有图片或视频亦包括在内)由作者上传并发布,文章内容仅代表作者本人观点,简书系信息发布平台,仅提供信息存储服务。