从零实现LLM项目：PyTorch模型权重加载的替代方案详解-优快云博客

从零实现LLM项目：PyTorch模型权重加载的替代方案详解

LLMs-from-scratch 从零开始逐步指导开发者构建自己的大型语言模型（LLM），旨在提供详细的步骤和原理说明，帮助用户深入理解并实践LLM的开发过程。项目地址: https://gitcode.com/gh_mirrors/ll/LLMs-from-scratch

引言

在构建大型语言模型(LLM)的过程中，模型权重的加载是一个关键步骤。本文将深入探讨在rasbt/LLMs-from-scratch项目中，如何通过PyTorch state dict和safetensors文件两种方式加载预训练权重，并分析它们的优缺点。

项目背景与模型配置

rasbt/LLMs-from-scratch项目旨在从零开始构建大型语言模型，其中第五章重点介绍了GPT模型的实现。项目提供了多种GPT模型变体的配置：

BASE_CONFIG = {
    "vocab_size": 50257,    # 词汇表大小
    "context_length": 1024, # 上下文长度
    "drop_rate": 0.0,       # 丢弃率
    "qkv_bias": True        # 查询-键-值偏置
}

model_configs = {
    "gpt2-small (124M)": {"emb_dim": 768, "n_layers": 12, "n_heads": 12},
    "gpt2-medium (355M)": {"emb_dim": 1024, "n_layers": 24, "n_heads": 16},
    "gpt2-large (774M)": {"emb_dim": 1280, "n_layers": 36, "n_heads": 20},
    "gpt2-xl (1558M)": {"emb_dim": 1600, "n_layers": 48, "n_heads": 25},
}

用户可以根据计算资源选择合适的模型规模，本文以"gpt2-small (124M)"为例进行说明。

传统PyTorch state dict权重加载

下载权重文件

项目提供了转换好的PyTorch权重文件，可以直接下载使用：

file_name = "gpt2-small-124M.pth"

import os
import urllib.request

url = "权重文件下载地址"  # 实际项目中应替换为真实URL

if not os.path.exists(file_name):
    urllib.request.urlretrieve(url, file_name)
    print(f"Downloaded to {file_name}")

加载权重到模型

加载过程非常简单直接：

import torch
from llms_from_scratch.ch04 import GPTModel

# 初始化模型
gpt = GPTModel(BASE_CONFIG)

# 加载权重
gpt.load_state_dict(torch.load(file_name, weights_only=True))

# 设置为评估模式
gpt.eval()

# 转移到合适的设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
gpt.to(device)

关键点说明：

weights_only=True参数是PyTorch 2.0+版本引入的安全特性，确保只加载权重数据而不执行潜在恶意代码
加载后必须调用eval()方法将模型设置为评估模式
根据硬件情况将模型转移到GPU或CPU

生成文本示例

加载完成后，可以使用模型生成文本：

import tiktoken
from llms_from_scratch.ch05 import generate, text_to_token_ids, token_ids_to_text

torch.manual_seed(123)  # 设置随机种子保证可重复性

tokenizer = tiktoken.get_encoding("gpt2")

token_ids = generate(
    model=gpt.to(device),
    idx=text_to_token_ids("Every effort moves", tokenizer).to(device),
    max_new_tokens=30,
    context_size=BASE_CONFIG["context_length"],
    top_k=1,
    temperature=1.0
)

print("Output text:\n", token_ids_to_text(token_ids, tokenizer))

使用safetensors文件加载权重

safetensors简介

safetensors是一种新兴的权重存储格式，相比传统PyTorch state dict有以下优势：

安全性：仅存储张量数据，不包含可执行代码
加载速度快：优化了序列化/反序列化过程
跨框架兼容性：支持多种深度学习框架

下载与加载

下载过程与state dict类似：

file_name = "gpt2-small-124M.safetensors"

if not os.path.exists(file_name):
    urllib.request.urlretrieve(url, file_name)
    print(f"Downloaded to {file_name}")

加载时需要使用专门的库：

from safetensors.torch import load_file

gpt = GPTModel(BASE_CONFIG)
gpt.load_state_dict(load_file(file_name))
gpt.eval()

生成文本验证

生成文本的代码与之前完全相同，可以用来验证两种加载方式的结果一致性。

技术对比与选择建议

| 特性 | PyTorch state dict | safetensors | |---------------------|--------------------|-------------| | 安全性 | 中等(需weights_only) | 高 | | 加载速度 | 一般 | 快 | | 框架支持 | PyTorch专用 | 多框架支持 | | 文件大小 | 较大 | 较小 | | 社区生态 | 成熟 | 新兴 |

选择建议：