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