File size: 3,828 Bytes
3e0ec22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | # SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
"""
Llama2-70B weights are saved as 8 sharded checkpoints. Loading weights for a
single layer is slow since we load all 80 layers into memory to construct the
model. This script repacks the weights into checkpoints chunked by layers to
speed up development.
"""
import argparse
import math
import shutil
from collections import defaultdict
from pathlib import Path
import torch
from tqdm import tqdm
def layer_num(key):
if "layers" in key:
return int(key.split("layers.")[1].split(".")[0])
return 0
def chunk_key(key, chunk_size):
"""
Return the chunk number that a key should go into
"""
chunk_id = layer_num(key) // chunk_size
print(f"Key: {key} -> chunk_id: {chunk_id}")
return chunk_id
def repack(in_dir, out_dir, chunk_size, num_layers, hidden_size):
"""
Repack llama2-70b weights into checkpoints chunked by layers.
Non-layer weights are saved in the first checkpoint.
"""
num_chunks = math.ceil(num_layers / chunk_size)
print(f"Repacking {num_layers} layers into {num_chunks} chunks of size {chunk_size}")
checkpoints = sorted(Path(in_dir).glob("*.pth"))
merged_checkpoints = defaultdict(list)
assert len(checkpoints) > 0, f"no checkpoint files found in {in_dir}"
print(f"Loading {len(checkpoints)} checkpoint files")
for ckpt in tqdm(checkpoints):
print(f"Checkpoint file: {ckpt}")
loaded_ckpt = torch.load(ckpt, map_location="cpu")
for key, value in loaded_ckpt.items():
merged_checkpoints[key].append(value)
# concat checkpoint values
chunks = [dict() for _ in range(num_chunks)]
for key, value in merged_checkpoints.items():
if len(value) == 1 or "norm" in key:
val = value[0]
else:
if (key == "tok_embeddings.weight" or key == "output.weight") and value[0].shape[1] == hidden_size:
# Concatenate along dimension 0 for llama3 token embeddings weight and lm head
val = torch.cat(value, dim=0)
else:
# cat_dim is index of the smallest dimension in value[0].shape
cat_dim = torch.argmin(torch.tensor(value[0].shape))
val = torch.cat(value, dim=cat_dim)
chunk_id = chunk_key(key, chunk_size)
chunks[chunk_id][key] = val
# save chunks and copy params.json if needed
out_dir = Path(out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
params_file = Path(in_dir) / "params.json"
if params_file.exists() and not (out_dir / "params.json").exists():
shutil.copy(params_file, out_dir)
print(f"Copied params.json to {out_dir}")
for i, chunk in enumerate(chunks):
# each chunk file name should tell which layers are in it
start_layer = i * chunk_size
end_layer = (i + 1) * chunk_size - 1
end_layer = min(end_layer, num_layers - 1)
out_file = out_dir / f"layers_{start_layer}-{end_layer}.pth"
torch.save(chunk, out_file)
print(f"Saved {out_file}")
if __name__ == "__main__":
# Take in command line arguments
parser = argparse.ArgumentParser(description="Repack llama2-70b weights")
parser.add_argument("in_dir", type=str, help="input directory")
parser.add_argument("out_dir", type=str, help="output directory")
parser.add_argument("chunk_size", type=int, default=10, help="number of layers per chunk")
parser.add_argument("-n", "--num_layers", type=int, default=80, help="total number of layers")
parser.add_argument("-hs", "--hidden_size", type=int, default=8192, help="hidden size of the model")
args = parser.parse_args()
repack(args.in_dir, args.out_dir, args.chunk_size, args.num_layers, args.hidden_size)
|