|
Supplementary code for the Build a Large Language Model From Scratch book by Sebastian Raschka Code repository: https://github.com/rasbt/LLMs-from-scratch |
|
Bonus Code for Chapter 5#
Alternative Weight Loading from Hugging Face Model Hub Via safetensors#
In the main chapter, we loaded the GPT model weights directly from OpenAI
This notebook provides alternative weight loading code to load the model weights from the Hugging Face Model Hub using
.safetensorsfilesThis is conceptually the same as loading weights of a PyTorch model from via the state-dict method described in chapter 5:
state_dict = torch.load("model_state_dict.pth")
model.load_state_dict(state_dict)
The appeal of
.safetensorsfiles lies in their secure design, as they only store tensor data and avoid the execution of potentially malicious code during loadingIn newer versions of PyTorch (e.g., 2.0 and newer), a
weights_only=Trueargument can be used withtorch.load(e.g.,torch.load("model_state_dict.pth", weights_only=True)) to improve safety by skipping the execution of code and loading only the weights (this is now enabled by default in PyTorch 2.6 and newer)
# pip install safetensors
from importlib.metadata import version
pkgs = ["numpy", "torch", "safetensors"]
for p in pkgs:
print(f"{p} version: {version(p)}")
numpy version: 2.2.3
---------------------------------------------------------------------------
PackageNotFoundError Traceback (most recent call last)
Cell In[2], line 5
3 pkgs = ["numpy", "torch", "safetensors"]
4 for p in pkgs:
----> 5 print(f"{p} version: {version(p)}")
File /Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/importlib/metadata/__init__.py:946, in version(distribution_name)
939 def version(distribution_name):
940 """Get the version string for the named package.
941
942 :param distribution_name: The name of the distribution package to query.
943 :return: The version string for the package as defined in the package's
944 "Version" metadata key.
945 """
--> 946 return distribution(distribution_name).version
File /Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/importlib/metadata/__init__.py:919, in distribution(distribution_name)
913 def distribution(distribution_name):
914 """Get the ``Distribution`` instance for the named package.
915
916 :param distribution_name: The name of the distribution package as a string.
917 :return: A ``Distribution`` instance (or subclass thereof).
918 """
--> 919 return Distribution.from_name(distribution_name)
File /Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/importlib/metadata/__init__.py:518, in Distribution.from_name(cls, name)
516 return dist
517 else:
--> 518 raise PackageNotFoundError(name)
PackageNotFoundError: No package metadata was found for torch
from llms_from_scratch.ch04 import GPTModel
# For llms_from_scratch installation instructions, see:
# https://github.com/rasbt/LLMs-from-scratch/tree/main/pkg
BASE_CONFIG = {
"vocab_size": 50257, # Vocabulary size
"context_length": 1024, # Context length
"drop_rate": 0.0, # Dropout rate
"qkv_bias": True # Query-key-value bias
}
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},
}
CHOOSE_MODEL = "gpt2-small (124M)"
BASE_CONFIG.update(model_configs[CHOOSE_MODEL])
import os
import urllib.request
from safetensors.torch import load_file
URL_DIR = {
"gpt2-small (124M)": "gpt2", # works ok
"gpt2-medium (355M)": "gpt2-medium", # this file seems to have issues via `generate`
"gpt2-large (774M)": "gpt2-large", # works ok
"gpt2-xl (1558M)": "gpt2-xl" # works ok
}
url = f"https://huggingface.co/openai-community/{URL_DIR[CHOOSE_MODEL]}/resolve/main/model.safetensors"
output_file = f"model-{URL_DIR[CHOOSE_MODEL]}.safetensors"
# Download file
if not os.path.exists(output_file):
urllib.request.urlretrieve(url, output_file)
# Load file
state_dict = load_file(output_file)
def assign(left, right):
if left.shape != right.shape:
raise ValueError(f"Shape mismatch. Left: {left.shape}, Right: {right.shape}")
return torch.nn.Parameter(right.detach())
def load_weights_into_gpt(gpt, params):
gpt.pos_emb.weight = assign(gpt.pos_emb.weight, params["wpe.weight"])
gpt.tok_emb.weight = assign(gpt.tok_emb.weight, params["wte.weight"])
for b in range(len(gpt.trf_blocks)):
q_w, k_w, v_w = torch.chunk(
params[f"h.{b}.attn.c_attn.weight"], 3, axis=-1)
gpt.trf_blocks[b].att.W_query.weight = assign(
gpt.trf_blocks[b].att.W_query.weight, q_w.T)
gpt.trf_blocks[b].att.W_key.weight = assign(
gpt.trf_blocks[b].att.W_key.weight, k_w.T)
gpt.trf_blocks[b].att.W_value.weight = assign(
gpt.trf_blocks[b].att.W_value.weight, v_w.T)
q_b, k_b, v_b = torch.chunk(
params[f"h.{b}.attn.c_attn.bias"], 3, axis=-1)
gpt.trf_blocks[b].att.W_query.bias = assign(
gpt.trf_blocks[b].att.W_query.bias, q_b)
gpt.trf_blocks[b].att.W_key.bias = assign(
gpt.trf_blocks[b].att.W_key.bias, k_b)
gpt.trf_blocks[b].att.W_value.bias = assign(
gpt.trf_blocks[b].att.W_value.bias, v_b)
gpt.trf_blocks[b].att.out_proj.weight = assign(
gpt.trf_blocks[b].att.out_proj.weight,
params[f"h.{b}.attn.c_proj.weight"].T)
gpt.trf_blocks[b].att.out_proj.bias = assign(
gpt.trf_blocks[b].att.out_proj.bias,
params[f"h.{b}.attn.c_proj.bias"])
gpt.trf_blocks[b].ff.layers[0].weight = assign(
gpt.trf_blocks[b].ff.layers[0].weight,
params[f"h.{b}.mlp.c_fc.weight"].T)
gpt.trf_blocks[b].ff.layers[0].bias = assign(
gpt.trf_blocks[b].ff.layers[0].bias,
params[f"h.{b}.mlp.c_fc.bias"])
gpt.trf_blocks[b].ff.layers[2].weight = assign(
gpt.trf_blocks[b].ff.layers[2].weight,
params[f"h.{b}.mlp.c_proj.weight"].T)
gpt.trf_blocks[b].ff.layers[2].bias = assign(
gpt.trf_blocks[b].ff.layers[2].bias,
params[f"h.{b}.mlp.c_proj.bias"])
gpt.trf_blocks[b].norm1.scale = assign(
gpt.trf_blocks[b].norm1.scale,
params[f"h.{b}.ln_1.weight"])
gpt.trf_blocks[b].norm1.shift = assign(
gpt.trf_blocks[b].norm1.shift,
params[f"h.{b}.ln_1.bias"])
gpt.trf_blocks[b].norm2.scale = assign(
gpt.trf_blocks[b].norm2.scale,
params[f"h.{b}.ln_2.weight"])
gpt.trf_blocks[b].norm2.shift = assign(
gpt.trf_blocks[b].norm2.shift,
params[f"h.{b}.ln_2.bias"])
gpt.final_norm.scale = assign(gpt.final_norm.scale, params["ln_f.weight"])
gpt.final_norm.shift = assign(gpt.final_norm.shift, params["ln_f.bias"])
gpt.out_head.weight = assign(gpt.out_head.weight, params["wte.weight"])
import torch
gpt = GPTModel(BASE_CONFIG)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
load_weights_into_gpt(gpt, state_dict)
gpt.to(device);
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))
Output text:
Every effort moves forward, but it's not enough.
"I'm not going to sit here and say, 'I'm not going to do this,'