问题1:如何运行7B以上的模型
13B的参数在一张A100上跑
- 合并00, 01两个pth文件
运行指令:python merge_weights.py --input_dir 00,01.pth存放的目录 --model_size 13B
合并后的参数文件放在同一个目录下
# Original copyright by Jason Phang
# https://github.com/zphang
# Taken here
# https://github.com/huggingface/transformers/pull/21955/commits/8978f28e6c44b083c0b190d3931902c2904c940a#diff-110a445233a8b15a0875998eeaf75cb8607b38a5daa736291dd058766879bbdd
import argparse
import json
import os
import shutil
import torch
"""
Sample usage:
```
python merge_weights.py --input_dir D:\Downloads\LLaMA --model_size 13B
```
"""
INTERMEDIATE_SIZE_MAP = {
"7B": 11008,
"13B": 13824,
"30B": 17920,
"65B": 22016,
}
NUM_SHARDS = {
"7B": 1,
"13B": 2,
"30B": 4,
"65B": 8,
}
def read_json(path):
with open(path, "r") as f:
return json.loads(f.read())
def write_model(input_base_path, model_size):
assert model_size in INTERMEDIATE_SIZE_MAP
params = read_json(os.path.join(input_base_path, "params.json"))
num_shards = NUM_SHARDS[model_size]
n_layers = params["n_layers"]
n_heads = params["n_heads"]
n_heads_per_shard = n_heads // num_shards
dim = params["dim"]
dims_per_head = dim // n_heads
# Load weights
if model_size == "7B":
loaded = torch.load(os.path.join(input_base_path, "consolidated.00.pth"), map_location="cpu")
else:
loaded = [
torch.load(os.path.join(input_base_path, f"consolidated.{i:02d}.pth"), map_location="cpu")
for i in range(num_shards)
]
state_dict = {}
for layer_i in range(n_layers):
if model_size == "7B":
state_dict |= {
f"layers.{layer_i}.attention.wq.weight": loaded[
f"layers.{layer_i}.attention.wq.weight"
],
f"layers.{layer_i}.attention.wk.weight": loaded[
f"layers.{layer_i}.attention.wk.weight"
],
f"layers.{layer_i}.attention.wv.weight": loaded[
f"layers.{layer_i}.attention.wv.weight"
],
f"layers.{layer_i}.attention.wo.weight": loaded[
f"layers.{layer_i}.attention.wo.weight"
],
f"layers.{layer_i}.feed_forward.w1.weight": loaded[
f"layers.{layer_i}.feed_forward.w1.weight"
],
f"layers.{layer_i}.feed_forward.w2.weight": loaded[
f"layers.{layer_i}.feed_forward.w2.weight"
],
f"layers.{layer_i}.feed_forward.w3.weight": loaded[
f"layers.{layer_i}.feed_forward.w3.weight"
],
f"layers.{layer_i}.attention_norm.weight": loaded[
f"layers.{layer_i}.attention_norm.weight"
],
f"layers.{layer_i}.ffn_norm.weight": loaded[f"layers.{layer_i}.ffn_norm.weight"],
}
else:
state_dict.update({
f"layers.{layer_i}.attention_norm.weight": loaded[0][
f"layers.{layer_i}.attention_norm.weight"
],
f"layers.{layer_i}.ffn_norm.weight": loaded[0][f"layers.{layer_i}.ffn_norm.weight"],
})
state_dict[f"layers.{layer_i}.attention.wq.weight"] = torch.cat(
[
loaded[i][f"layers.{layer_i}.attention.wq.weight"].view(n_heads_per_shard, dims_per_head, dim)
for i in range(num_shards)
],
dim=0,
).reshape(dim, dim)
state_dict[f"layers.{layer_i}.attention.wk.weight"] = torch.cat(
[
loaded[i][f"layers.{layer_i}.attention.wk.weight"].view(n_heads_per_shard, dims_per_head, dim)
for i in range(num_shards)
],
dim=0,
).reshape(dim, dim)
state_dict[f"layers.{layer_i}.attention.wv.weight"] = torch.cat(
[
loaded[i][f"layers.{layer_i}.attention.wv.weight"].view(n_heads_per_shard, dims_per_head, dim)
for i in range(num_shards)
],
dim=0,
).reshape(dim, dim)
state_dict[f"layers.{layer_i}.attention.wo.weight"] = torch.cat(
[loaded[i][f"layers.{layer_i}.attention.wo.weight"] for i in range(num_shards)], dim=1
)
state_dict[f"layers.{layer_i}.feed_forward.w1.weight"] = torch.cat(
[loaded[i][f"layers.{layer_i}.feed_forward.w1.weight"] for i in range(num_shards)], dim=0
)
state_dict[f"layers.{layer_i}.feed_forward.w2.weight"] = torch.cat(
[loaded[i][f"layers.{layer_i}.feed_forward.w2.weight"] for i in range(num_shards)], dim=1
)
state_dict[f"layers.{layer_i}.feed_forward.w3.weight"] = torch.cat(
[loaded[i][f"layers.{layer_i}.feed_forward.w3.weight"] for i in range(num_shards)], dim=0
)
if model_size == "7B":
state_dict |= {
"tok_embeddings.weight": loaded["tok_embeddings.weight"],
"norm.weight": loaded["norm.weight"],
"output.weight": loaded["output.weight"],
}
else:
state_dict.update ({
"norm.weight": loaded[0]["norm.weight"],
"tok_embeddings.weight": torch.cat(
[loaded[i]["tok_embeddings.weight"] for i in range(num_shards)], dim=1
),
"output.weight": torch.cat([loaded[i]["output.weight"] for i in range(num_shards)], dim=0),
})
torch.save(state_dict, 'merged.pth')
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
"--input_dir",
help="Location of LLaMA weights, which contains tokenizer.model and model folders",
)
parser.add_argument(
"--model_size",
choices=["7B", "13B", "30B", "65B"],
)
args = parser.parse_args()
write_model(
input_base_path=os.path.join(args.input_dir, args.model_size),
model_size=args.model_size,
)
if __name__ == "__main__":
main()
之后用torchrun --nproc_per_node 1 example --ckpt_dir 合并后的参数文件所在的目录 --tokenizer_path tokenizer.model
运行即可。
2. 用两张卡跑
待补充
问题2:相同的prompt输入到llama中,不管几次,都会得到相同的answer
由于llama除了7B以外都是多GPU并行,所以在为了确保并行计算在每个GPU上的可复现性,会固定随机数生成器的种子,即torch.manual_seed(1)
。
example.py -> setup_model_parallel()
def setup_model_parallel() -> Tuple[int, int]:
local_rank = int(os.environ.get("LOCAL_RANK", -1))
world_size = int(os.environ.get("WORLD_SIZE", -1))
torch.distributed.init_process_group("nccl")
initialize_model_parallel(world_size)
torch.cuda.set_device(local_rank)
# seed must be the same in all processes
torch.manual_seed(1)
return local_rank, world_size