Skip to content

GitLab

Commit 63c300ba authored by wxj's avatar wxj
Browse files

Update loader_llama_mistral.py

parent be4dda7b
Pipeline #2651 passed with stage
Hide whitespace changes
Inline Side-by-side
Showing with 659 additions and 659 deletions
tools/checkpoint/loader_llama_mistral.py
View file @ 63c300ba
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
import json import json
import os import os
import sys import sys
import torch import torch
try: try:
import transformers import transformers
except ImportError: except ImportError:
raise ImportError("The 'transformers' package is not installed.") raise ImportError("The 'transformers' package is not installed.")
import gc import gc
import shutil import shutil
from tqdm import tqdm from tqdm import tqdm
import types import types
def add_arguments(parser): def add_arguments(parser):
group = parser.add_argument_group(title='Llama/Mistral loader.') group = parser.add_argument_group(title='Llama/Mistral loader.')
# TODO(jbarker): Need assertion to make sure *exactly* one of these is used # TODO(jbarker): Need assertion to make sure *exactly* one of these is used
parser.add_argument('--model-size', type=str, required=True, parser.add_argument('--model-size', type=str, required=True,
choices=['llama2-7B', 'llama2-13B', 'llama2-70B', 'llama2-7Bf', 'llama2-13Bf', 'llama2-70Bf', 'llama3', 'mistral', 'yi-34B', 'qwen2.5'], choices=['llama2-7B', 'llama2-13B', 'llama2-70B', 'llama2-7Bf', 'llama2-13Bf', 'llama2-70Bf', 'llama3', 'mistral', 'yi-34B', 'qwen2.5'],
help='Select model size/type') help='Select model size/type')
parser.add_argument('--checkpoint-type', type=str, required=True, parser.add_argument('--checkpoint-type', type=str, required=True,
choices=['meta', 'hf'], choices=['meta', 'hf'],
help='Type of checkpoint to convert, options are "meta" or "hf"') help='Type of checkpoint to convert, options are "meta" or "hf"')
parser.add_argument('--bf16', action='store_true', help='Whether to load weights in bf16.') parser.add_argument('--bf16', action='store_true', help='Whether to load weights in bf16.')
parser.add_argument('--fp16', action='store_true', help='Whether to load weights in fp16.') parser.add_argument('--fp16', action='store_true', help='Whether to load weights in fp16.')
group.add_argument('--true-vocab-size', type=int, default=None, group.add_argument('--true-vocab-size', type=int, default=None,
help='original size of vocab, if specified will trim padding from embedding table.') help='original size of vocab, if specified will trim padding from embedding table.')
group.add_argument('--vocab-file', type=str, default=None, group.add_argument('--vocab-file', type=str, default=None,
help='Path to the vocab file. If specified will use this to get vocab size and ' help='Path to the vocab file. If specified will use this to get vocab size and '
'trim padding from the embedding table.') 'trim padding from the embedding table.')
group.add_argument('--tokenizer-model', required=True, group.add_argument('--tokenizer-model', required=True,
help='Tokenizer model file.') help='Tokenizer model file.')
group.add_argument('--megatron-path', type=str, default=None, group.add_argument('--megatron-path', type=str, default=None,
help='Base directory of Megatron repository') help='Base directory of Megatron repository')
group.add_argument("--make-vocab-size-divisible-by", type=int, default=None, help="Make vocab size divisible by") group.add_argument("--make-vocab-size-divisible-by", type=int, default=None, help="Make vocab size divisible by")
group.add_argument('--loader-transformer-impl', default='local', group.add_argument('--loader-transformer-impl', default='local',
choices=['local', 'transformer_engine'], choices=['local', 'transformer_engine'],
help='Which Transformer implementation to use.') help='Which Transformer implementation to use.')
def verify_transformers_version(): def verify_transformers_version():
major, minor, patch = map(int, transformers.__version__.split('.')) major, minor, patch = map(int, transformers.__version__.split('.'))
assert major >= 4 and minor >= 31 assert major >= 4 and minor >= 31
NUM_SHARDS = { NUM_SHARDS = {
"llama2-7B": 1, "llama2-7B": 1,
"llama2-7Bf": 1, "llama2-7Bf": 1,
"llama2-13B": 2, "llama2-13B": 2,
"llama2-13Bf": 2, "llama2-13Bf": 2,
"llama2-70B": 8, "llama2-70B": 8,
"llama2-70Bf": 8, "llama2-70Bf": 8,
} }
def compute_intermediate_size(n, ffn_dim_multiplier=1, multiple_of=256): def compute_intermediate_size(n, ffn_dim_multiplier=1, multiple_of=256):
return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3)) + multiple_of - 1) // multiple_of) return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3)) + multiple_of - 1) // multiple_of)
def read_json(path): def read_json(path):
with open(path, "r") as f: with open(path, "r") as f:
return json.load(f) return json.load(f)
def write_json(text, path): def write_json(text, path):
with open(path, "w") as f: with open(path, "w") as f:
json.dump(text, f) json.dump(text, f)
# This conversion is adapted from # This conversion is adapted from
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/convert_llama_weights_to_hf.py # https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/convert_llama_weights_to_hf.py
def convert_to_hf(model_path, input_base_path, model_size, tokenizer_path): def convert_to_hf(model_path, input_base_path, model_size, tokenizer_path):
if "llama2" in model_size: if "llama2" in model_size:
from transformers import LlamaConfig as ModelConfig from transformers import LlamaConfig as ModelConfig
from transformers import LlamaTokenizer, LlamaTokenizerFast from transformers import LlamaTokenizer, LlamaTokenizerFast
else: else:
raise NotImplementedError(f"converting {model_size} is only supported using HuggingFace weights") raise NotImplementedError(f"converting {model_size} is only supported using HuggingFace weights")
# for backward compatibility, before you needed the repo to be called `my_repo/model_size` # for backward compatibility, before you needed the repo to be called `my_repo/model_size`
if not os.path.isfile(os.path.join(input_base_path, "params.json")): if not os.path.isfile(os.path.join(input_base_path, "params.json")):
input_base_path = os.path.join(input_base_path, model_size) input_base_path = os.path.join(input_base_path, model_size)
os.makedirs(model_path, exist_ok=True) os.makedirs(model_path, exist_ok=True)
params = read_json(os.path.join(input_base_path, "params.json")) params = read_json(os.path.join(input_base_path, "params.json"))
num_shards = NUM_SHARDS[model_size] num_shards = NUM_SHARDS[model_size]
params = params.get("model", params) params = params.get("model", params)
n_layers = params["n_layers"] n_layers = params["n_layers"]
n_heads = params["n_heads"] n_heads = params["n_heads"]
n_heads_per_shard = n_heads // num_shards n_heads_per_shard = n_heads // num_shards
dim = params["dim"] dim = params["dim"]
dims_per_head = dim // n_heads dims_per_head = dim // n_heads
base = params.get("rope_theta", 10000.0) base = params.get("rope_theta", 10000.0)
inv_freq = 1.0 / (base ** (torch.arange(0, dims_per_head, 2).float() / dims_per_head)) inv_freq = 1.0 / (base ** (torch.arange(0, dims_per_head, 2).float() / dims_per_head))
if base > 10000.0: if base > 10000.0:
max_position_embeddings = 32768 if "mistral" in model_size else 16384 max_position_embeddings = 32768 if "mistral" in model_size else 16384
else: else:
max_position_embeddings = 4096 max_position_embeddings = 4096
if "llama2" in model_size: if "llama2" in model_size:
tokenizer_class = LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast tokenizer_class = LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
else: else:
raise AttributeError(f"model_size={model_size} not supported") raise AttributeError(f"model_size={model_size} not supported")
if tokenizer_path is not None: if tokenizer_path is not None:
if "llama2" in model_size: if "llama2" in model_size:
tokenizer = tokenizer_class(tokenizer_path) tokenizer = tokenizer_class(tokenizer_path)
tokenizer.save_pretrained(model_path) tokenizer.save_pretrained(model_path)
vocab_size = tokenizer.vocab_size if tokenizer_path is not None else 32000 vocab_size = tokenizer.vocab_size if tokenizer_path is not None else 32000
else: else:
raise AttributeError(f"model_size={model_size} is not supported") raise AttributeError(f"model_size={model_size} is not supported")
if params.get("n_kv_heads", None) is not None: if params.get("n_kv_heads", None) is not None:
num_key_value_heads = params["n_kv_heads"] # for GQA / MQA num_key_value_heads = params["n_kv_heads"] # for GQA / MQA
num_local_key_value_heads = n_heads_per_shard // num_key_value_heads num_local_key_value_heads = n_heads_per_shard // num_key_value_heads
key_value_dim = dim // num_key_value_heads key_value_dim = dim // num_key_value_heads
else: # compatibility with other checkpoints else: # compatibility with other checkpoints
num_key_value_heads = n_heads num_key_value_heads = n_heads
num_local_key_value_heads = n_heads_per_shard num_local_key_value_heads = n_heads_per_shard
key_value_dim = dim key_value_dim = dim
# permute for sliced rotary # permute for sliced rotary
def permute(w, n_heads=n_heads, dim1=dim, dim2=dim): def permute(w, n_heads=n_heads, dim1=dim, dim2=dim):
return w.view(n_heads, dim1 // n_heads // 2, 2, dim2).transpose(1, 2).reshape(dim1, dim2) return w.view(n_heads, dim1 // n_heads // 2, 2, dim2).transpose(1, 2).reshape(dim1, dim2)
print(f"Fetching all parameters from the checkpoint at {input_base_path}.") print(f"Fetching all parameters from the checkpoint at {input_base_path}.")
# Load weights # Load weights
if num_shards == 1: if num_shards == 1:
# Not sharded # Not sharded
# (The sharded implementation would also work, but this is simpler.) # (The sharded implementation would also work, but this is simpler.)
loaded = torch.load(os.path.join(input_base_path, "consolidated.00.pth"), map_location="cpu") loaded = torch.load(os.path.join(input_base_path, "consolidated.00.pth"), map_location="cpu")
else: else:
# Sharded # Sharded
loaded = [ loaded = [
torch.load(os.path.join(input_base_path, f"consolidated.{i:02d}.pth"), map_location="cpu") torch.load(os.path.join(input_base_path, f"consolidated.{i:02d}.pth"), map_location="cpu")
for i in range(num_shards) for i in range(num_shards)
] ]
param_count = 0 param_count = 0
index_dict = {"weight_map": {}} index_dict = {"weight_map": {}}
for layer_i in range(n_layers): for layer_i in range(n_layers):
filename = f"pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin" filename = f"pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin"
if num_shards == 1: if num_shards == 1:
# Unsharded # Unsharded
q_proj = loaded[f"layers.{layer_i}.attention.wq.weight"] q_proj = loaded[f"layers.{layer_i}.attention.wq.weight"]
k_proj = loaded[f"layers.{layer_i}.attention.wk.weight"] k_proj = loaded[f"layers.{layer_i}.attention.wk.weight"]
if ("llama2" in model_size) or ("mistral" in model_size): if ("llama2" in model_size) or ("mistral" in model_size):
q_proj = permute(q_proj) q_proj = permute(q_proj)
k_proj = permute(k_proj) k_proj = permute(k_proj)
state_dict = { state_dict = {
f"model.layers.{layer_i}.self_attn.q_proj.weight": q_proj, f"model.layers.{layer_i}.self_attn.q_proj.weight": q_proj,
f"model.layers.{layer_i}.self_attn.k_proj.weight": k_proj, f"model.layers.{layer_i}.self_attn.k_proj.weight": k_proj,
f"model.layers.{layer_i}.self_attn.v_proj.weight": loaded[f"layers.{layer_i}.attention.wv.weight"], f"model.layers.{layer_i}.self_attn.v_proj.weight": loaded[f"layers.{layer_i}.attention.wv.weight"],
f"model.layers.{layer_i}.self_attn.o_proj.weight": loaded[f"layers.{layer_i}.attention.wo.weight"], f"model.layers.{layer_i}.self_attn.o_proj.weight": loaded[f"layers.{layer_i}.attention.wo.weight"],
f"model.layers.{layer_i}.mlp.gate_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w1.weight"], f"model.layers.{layer_i}.mlp.gate_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w1.weight"],
f"model.layers.{layer_i}.mlp.down_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w2.weight"], f"model.layers.{layer_i}.mlp.down_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w2.weight"],
f"model.layers.{layer_i}.mlp.up_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w3.weight"], f"model.layers.{layer_i}.mlp.up_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w3.weight"],
f"model.layers.{layer_i}.input_layernorm.weight": loaded[f"layers.{layer_i}.attention_norm.weight"], f"model.layers.{layer_i}.input_layernorm.weight": loaded[f"layers.{layer_i}.attention_norm.weight"],
f"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[f"layers.{layer_i}.ffn_norm.weight"], f"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[f"layers.{layer_i}.ffn_norm.weight"],
} }
else: else:
# Sharded # Sharded
# Note that attention.w{q,k,v,o}, feed_fordward.w[1,2,3], attention_norm.weight and ffn_norm.weight share # Note that attention.w{q,k,v,o}, feed_fordward.w[1,2,3], attention_norm.weight and ffn_norm.weight share
# the same storage object, saving attention_norm and ffn_norm will save other weights too, which is # the same storage object, saving attention_norm and ffn_norm will save other weights too, which is
# redundant as other weights will be stitched from multiple shards. To avoid that, they are cloned. # redundant as other weights will be stitched from multiple shards. To avoid that, they are cloned.
state_dict = { state_dict = {
f"model.layers.{layer_i}.input_layernorm.weight": loaded[0][ f"model.layers.{layer_i}.input_layernorm.weight": loaded[0][
f"layers.{layer_i}.attention_norm.weight" f"layers.{layer_i}.attention_norm.weight"
].clone(), ].clone(),
f"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[0][ f"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[0][
f"layers.{layer_i}.ffn_norm.weight" f"layers.{layer_i}.ffn_norm.weight"
].clone(), ].clone(),
} }
state_dict[f"model.layers.{layer_i}.self_attn.q_proj.weight"] = permute( state_dict[f"model.layers.{layer_i}.self_attn.q_proj.weight"] = permute(
torch.cat( torch.cat(
[ [
loaded[i][f"layers.{layer_i}.attention.wq.weight"].view(n_heads_per_shard, dims_per_head, dim) loaded[i][f"layers.{layer_i}.attention.wq.weight"].view(n_heads_per_shard, dims_per_head, dim)
for i in range(num_shards) for i in range(num_shards)
], ],
dim=0, dim=0,
).reshape(dim, dim) ).reshape(dim, dim)
) )
state_dict[f"model.layers.{layer_i}.self_attn.k_proj.weight"] = permute( state_dict[f"model.layers.{layer_i}.self_attn.k_proj.weight"] = permute(
torch.cat( torch.cat(
[ [
loaded[i][f"layers.{layer_i}.attention.wk.weight"].view( loaded[i][f"layers.{layer_i}.attention.wk.weight"].view(
num_local_key_value_heads, dims_per_head, dim num_local_key_value_heads, dims_per_head, dim
) )
for i in range(num_shards) for i in range(num_shards)
], ],
dim=0, dim=0,
).reshape(key_value_dim, dim), ).reshape(key_value_dim, dim),
num_key_value_heads, num_key_value_heads,
key_value_dim, key_value_dim,
dim, dim,
) )
state_dict[f"model.layers.{layer_i}.self_attn.v_proj.weight"] = torch.cat( state_dict[f"model.layers.{layer_i}.self_attn.v_proj.weight"] = torch.cat(
[ [
loaded[i][f"layers.{layer_i}.attention.wv.weight"].view( loaded[i][f"layers.{layer_i}.attention.wv.weight"].view(
num_local_key_value_heads, dims_per_head, dim num_local_key_value_heads, dims_per_head, dim
) )
for i in range(num_shards) for i in range(num_shards)
], ],
dim=0, dim=0,
).reshape(key_value_dim, dim) ).reshape(key_value_dim, dim)
state_dict[f"model.layers.{layer_i}.self_attn.o_proj.weight"] = torch.cat( state_dict[f"model.layers.{layer_i}.self_attn.o_proj.weight"] = torch.cat(
[loaded[i][f"layers.{layer_i}.attention.wo.weight"] for i in range(num_shards)], dim=1 [loaded[i][f"layers.{layer_i}.attention.wo.weight"] for i in range(num_shards)], dim=1
) )
state_dict[f"model.layers.{layer_i}.mlp.gate_proj.weight"] = torch.cat( state_dict[f"model.layers.{layer_i}.mlp.gate_proj.weight"] = torch.cat(
[loaded[i][f"layers.{layer_i}.feed_forward.w1.weight"] for i in range(num_shards)], dim=0 [loaded[i][f"layers.{layer_i}.feed_forward.w1.weight"] for i in range(num_shards)], dim=0
) )
state_dict[f"model.layers.{layer_i}.mlp.down_proj.weight"] = torch.cat( state_dict[f"model.layers.{layer_i}.mlp.down_proj.weight"] = torch.cat(
[loaded[i][f"layers.{layer_i}.feed_forward.w2.weight"] for i in range(num_shards)], dim=1 [loaded[i][f"layers.{layer_i}.feed_forward.w2.weight"] for i in range(num_shards)], dim=1
) )
state_dict[f"model.layers.{layer_i}.mlp.up_proj.weight"] = torch.cat( state_dict[f"model.layers.{layer_i}.mlp.up_proj.weight"] = torch.cat(
[loaded[i][f"layers.{layer_i}.feed_forward.w3.weight"] for i in range(num_shards)], dim=0 [loaded[i][f"layers.{layer_i}.feed_forward.w3.weight"] for i in range(num_shards)], dim=0
) )
state_dict[f"model.layers.{layer_i}.self_attn.rotary_emb.inv_freq"] = inv_freq state_dict[f"model.layers.{layer_i}.self_attn.rotary_emb.inv_freq"] = inv_freq
for k, v in state_dict.items(): for k, v in state_dict.items():
index_dict["weight_map"][k] = filename index_dict["weight_map"][k] = filename
param_count += v.numel() param_count += v.numel()
torch.save(state_dict, os.path.join(model_path, filename)) torch.save(state_dict, os.path.join(model_path, filename))
filename = f"pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin" filename = f"pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin"
if num_shards == 1: if num_shards == 1:
# Unsharded # Unsharded
state_dict = { state_dict = {
"model.embed_tokens.weight": loaded["tok_embeddings.weight"], "model.embed_tokens.weight": loaded["tok_embeddings.weight"],
"model.norm.weight": loaded["norm.weight"], "model.norm.weight": loaded["norm.weight"],
"lm_head.weight": loaded["output.weight"], "lm_head.weight": loaded["output.weight"],
} }
else: else:
d = 0 if "llama3" in model_size else 1 d = 0 if "llama3" in model_size else 1
state_dict = { state_dict = {
"model.norm.weight": loaded[0]["norm.weight"], "model.norm.weight": loaded[0]["norm.weight"],
"model.embed_tokens.weight": torch.cat( "model.embed_tokens.weight": torch.cat(
[loaded[i]["tok_embeddings.weight"] for i in range(num_shards)], dim=d [loaded[i]["tok_embeddings.weight"] for i in range(num_shards)], dim=d
), ),
"lm_head.weight": torch.cat([loaded[i]["output.weight"] for i in range(num_shards)], dim=0), "lm_head.weight": torch.cat([loaded[i]["output.weight"] for i in range(num_shards)], dim=0),
} }
for k, v in state_dict.items(): for k, v in state_dict.items():
index_dict["weight_map"][k] = filename index_dict["weight_map"][k] = filename
param_count += v.numel() param_count += v.numel()
torch.save(state_dict, os.path.join(model_path, filename)) torch.save(state_dict, os.path.join(model_path, filename))
# Write configs # Write configs
index_dict["metadata"] = {"total_size": param_count * 2} index_dict["metadata"] = {"total_size": param_count * 2}
write_json(index_dict, os.path.join(model_path, "pytorch_model.bin.index.json")) write_json(index_dict, os.path.join(model_path, "pytorch_model.bin.index.json"))
ffn_dim_multiplier = params["ffn_dim_multiplier"] if "ffn_dim_multiplier" in params else 1 ffn_dim_multiplier = params["ffn_dim_multiplier"] if "ffn_dim_multiplier" in params else 1
multiple_of = params["multiple_of"] if "multiple_of" in params else 256 multiple_of = params["multiple_of"] if "multiple_of" in params else 256
config = ModelConfig( config = ModelConfig(
hidden_size=dim, hidden_size=dim,
intermediate_size=compute_intermediate_size(dim, ffn_dim_multiplier, multiple_of), intermediate_size=compute_intermediate_size(dim, ffn_dim_multiplier, multiple_of),
num_attention_heads=params["n_heads"], num_attention_heads=params["n_heads"],
num_hidden_layers=params["n_layers"], num_hidden_layers=params["n_layers"],
rms_norm_eps=params["norm_eps"], rms_norm_eps=params["norm_eps"],
num_key_value_heads=num_key_value_heads, num_key_value_heads=num_key_value_heads,
vocab_size=vocab_size, vocab_size=vocab_size,
rope_theta=base, rope_theta=base,
max_position_embeddings=max_position_embeddings, max_position_embeddings=max_position_embeddings,
) )
config.save_pretrained(model_path) config.save_pretrained(model_path)
# Make space so we can load the model properly now. # Make space so we can load the model properly now.
del state_dict del state_dict
del loaded del loaded
gc.collect() gc.collect()
return model_path return model_path
def load_args_from_checkpoint(args, model_size): def load_args_from_checkpoint(args, model_size):
# Read Llama args. # Read Llama args.
model_args_path = os.path.join(args.load, "config.json") model_args_path = os.path.join(args.load, "config.json")
with open(model_args_path) as f: with open(model_args_path) as f:
model_args = json.load(f) model_args = json.load(f)
# Update Megatron args. # Update Megatron args.
args.seq_length = 4096 args.seq_length = 4096
if "llama2" in model_size: if "llama2" in model_size:
# Correct bug in earlier conversion script. # Correct bug in earlier conversion script.
args.max_position_embeddings = 4096 args.max_position_embeddings = 4096
else: else:
args.max_position_embeddings = model_args["max_position_embeddings"] args.max_position_embeddings = model_args["max_position_embeddings"]
args.hidden_size = model_args["hidden_size"] args.hidden_size = model_args["hidden_size"]
args.num_attention_heads = model_args["num_attention_heads"] args.num_attention_heads = model_args["num_attention_heads"]
args.num_layers = model_args["num_hidden_layers"] args.num_layers = model_args["num_hidden_layers"]
args.global_batch_size = 1024 args.global_batch_size = 1024
args.norm_epsilon = model_args["rms_norm_eps"] args.norm_epsilon = model_args["rms_norm_eps"]
args.iteration = 1 # '0', 'release' don't work args.iteration = 1 # '0', 'release' don't work
args.position_embedding_type = "rope" args.position_embedding_type = "rope"
args.swiglu = True args.swiglu = True
args.normalization = "RMSNorm" args.normalization = "RMSNorm"
args.add_bias_linear = False args.add_bias_linear = False
args.untie_embeddings_and_output_weights = not model_args.get("tie_word_embeddings", False) args.untie_embeddings_and_output_weights = not model_args.get("tie_word_embeddings", False)
args.vocab_size = model_args["vocab_size"] args.vocab_size = model_args["vocab_size"]
args.padded_vocab_size = model_args["vocab_size"] args.padded_vocab_size = model_args["vocab_size"]
args.ffn_hidden_size = model_args["intermediate_size"] args.ffn_hidden_size = model_args["intermediate_size"]
if "num_key_value_heads" in model_args: if "num_key_value_heads" in model_args:
args.group_query_attention = True args.group_query_attention = True
args.num_query_groups = model_args["num_key_value_heads"] args.num_query_groups = model_args["num_key_value_heads"]
def set_preprocess_state(args, model, hf_model): def set_preprocess_state(args, model, hf_model):
'''Set embedding params.''' '''Set embedding params.'''
model.language_model.embedding.word_embeddings.weight.data.copy_( model.language_model.embedding.word_embeddings.weight.data.copy_(
hf_model.model.embed_tokens.weight) hf_model.model.embed_tokens.weight)
def set_postprocess_state(args, model, hf_model): def set_postprocess_state(args, model, hf_model):
'''Set output layer & norm params.''' '''Set output layer & norm params.'''
model.language_model.encoder.final_norm.weight.data.copy_(hf_model.model.norm.weight) model.language_model.encoder.final_norm.weight.data.copy_(hf_model.model.norm.weight)
if args.untie_embeddings_and_output_weights: if args.untie_embeddings_and_output_weights:
model.language_model.output_layer.weight.data.copy_(hf_model.lm_head.weight) model.language_model.output_layer.weight.data.copy_(hf_model.lm_head.weight)
def set_attn_state(args, layer, hf_layer): def set_attn_state(args, layer, hf_layer):
'''Set self-attention params.''' '''Set self-attention params.'''
# Get attention layer & state. # Get attention layer & state.
attn = layer.self_attention attn = layer.self_attention
hf_attn = hf_layer.self_attn hf_attn = hf_layer.self_attn
# Reshape loaded weights. # Reshape loaded weights.
tp = args.tensor_model_parallel_size tp = args.tensor_model_parallel_size
nh = args.num_attention_heads // tp nh = args.num_attention_heads // tp
ng = (args.num_query_groups if args.group_query_attention \ ng = (args.num_query_groups if args.group_query_attention \
else args.num_attention_heads) // tp else args.num_attention_heads) // tp
dim = args.kv_channels dim = args.kv_channels
assert nh % ng == 0 assert nh % ng == 0
# Copy weights (re-order dimensions for Megatron). # Copy weights (re-order dimensions for Megatron).
attn.query_key_value.weight.data.copy_(torch.cat([ attn.query_key_value.weight.data.copy_(torch.cat([
hf_attn.q_proj.weight.reshape((ng, dim*nh//ng, -1)), hf_attn.q_proj.weight.reshape((ng, dim*nh//ng, -1)),
hf_attn.k_proj.weight.reshape((ng, dim, -1)), hf_attn.k_proj.weight.reshape((ng, dim, -1)),
hf_attn.v_proj.weight.reshape((ng, dim, -1)), hf_attn.v_proj.weight.reshape((ng, dim, -1)),
], dim=1).reshape((-1, args.hidden_size))) ], dim=1).reshape((-1, args.hidden_size)))
if args.add_qkv_bias: if args.add_qkv_bias:
attn.query_key_value.bias.data.copy_(torch.cat([ attn.query_key_value.bias.data.copy_(torch.cat([
hf_attn.q_proj.bias.reshape((ng, dim*nh//ng)), hf_attn.q_proj.bias.reshape((ng, dim*nh//ng)),
hf_attn.k_proj.bias.reshape((ng, dim)), hf_attn.k_proj.bias.reshape((ng, dim)),
hf_attn.v_proj.bias.reshape((ng, dim)), hf_attn.v_proj.bias.reshape((ng, dim)),
], dim=1).reshape(-1)) ], dim=1).reshape(-1))
attn.dense.weight.data.copy_(hf_attn.o_proj.weight) attn.dense.weight.data.copy_(hf_attn.o_proj.weight)
def set_mlp_state(args, layer, hf_layer): def set_mlp_state(args, layer, hf_layer):
'''Set MLP params.''' '''Set MLP params.'''
mlp = layer.mlp mlp = layer.mlp
hf_mlp = hf_layer.mlp hf_mlp = hf_layer.mlp
mlp.dense_h_to_4h.weight.data.copy_(torch.cat([ mlp.dense_h_to_4h.weight.data.copy_(torch.cat([
hf_mlp.gate_proj.weight, hf_mlp.gate_proj.weight,
hf_mlp.up_proj.weight, hf_mlp.up_proj.weight,
], dim=0)) ], dim=0))
mlp.dense_4h_to_h.weight.data.copy_(hf_mlp.down_proj.weight) mlp.dense_4h_to_h.weight.data.copy_(hf_mlp.down_proj.weight)
def set_layer_state(args, model, hf_model, layer_idx): def set_layer_state(args, model, hf_model, layer_idx):
'''Set transformer layer params.''' '''Set transformer layer params.'''
layer = model.language_model.encoder.layers[layer_idx] layer = model.language_model.encoder.layers[layer_idx]
hf_layer = hf_model.model.layers[layer_idx] hf_layer = hf_model.model.layers[layer_idx]
set_attn_state(args, layer, hf_layer) set_attn_state(args, layer, hf_layer)
set_mlp_state(args, layer, hf_layer) set_mlp_state(args, layer, hf_layer)
layer.input_norm.weight.data.copy_(hf_layer.input_layernorm.weight) layer.input_norm.weight.data.copy_(hf_layer.input_layernorm.weight)
layer.post_attention_norm.weight.data.copy_(hf_layer.post_attention_layernorm.weight) layer.post_attention_norm.weight.data.copy_(hf_layer.post_attention_layernorm.weight)
def load_checkpoint_to_model(args): def load_checkpoint_to_model(args):
'''Set model params.''' '''Set model params.'''
from pretrain_gpt import model_provider from pretrain_gpt import model_provider
from transformers import AutoModelForCausalLM from transformers import AutoModelForCausalLM
# Load Huggingface model. # Load Huggingface model.
hf_model = AutoModelForCausalLM.from_pretrained(args.load, torch_dtype=args.params_dtype, low_cpu_mem_usage=True, device_map="cpu") hf_model = AutoModelForCausalLM.from_pretrained(args.load, torch_dtype=args.params_dtype, low_cpu_mem_usage=True, device_map="cpu")
# Init Megatron model. # Init Megatron model.
model = model_provider(True, True).to(args.params_dtype) model = model_provider(True, True).to(args.params_dtype)
# Set model state. # Set model state.
set_preprocess_state(args, model, hf_model) set_preprocess_state(args, model, hf_model)
set_postprocess_state(args, model, hf_model) set_postprocess_state(args, model, hf_model)
for layer_idx in tqdm(range(args.num_layers), "set layer states"): for layer_idx in tqdm(range(args.num_layers), "set layer states"):
set_layer_state(args, model, hf_model, layer_idx) set_layer_state(args, model, hf_model, layer_idx)
return model return model
def _load_checkpoint(queue, args): def _load_checkpoint(queue, args):
verify_transformers_version() verify_transformers_version()
# Search in directory above this. # Search in directory above this.
sys.path.append(os.path.abspath( sys.path.append(os.path.abspath(
os.path.join(os.path.dirname(__file__), os.path.join(os.path.dirname(__file__),
os.path.pardir, os.path.pardir,
os.path.pardir))) os.path.pardir)))
if args.megatron_path is not None: if args.megatron_path is not None:
sys.path.insert(0, args.megatron_path) sys.path.insert(0, args.megatron_path)
# Convert Meta checkpoint to HF format as an intermediate step # Convert Meta checkpoint to HF format as an intermediate step
if args.checkpoint_type == "meta": if args.checkpoint_type == "meta":
model_tmp_path = convert_to_hf(model_path=os.path.join(args.save_dir, 'tmp'), input_base_path=args.load_dir, model_size=args.model_size, tokenizer_path=args.tokenizer_model) model_tmp_path = convert_to_hf(model_path=os.path.join(args.save_dir, 'tmp'), input_base_path=args.load_dir, model_size=args.model_size, tokenizer_path=args.tokenizer_model)
args.load_dir = model_tmp_path args.load_dir = model_tmp_path
args.tokenizer_model = model_tmp_path # point to HF tokenizer model args.tokenizer_model = model_tmp_path # point to HF tokenizer model
try: try:
from megatron.training.arguments import parse_args, validate_args from megatron.training.arguments import parse_args, validate_args
from megatron.training.global_vars import set_args, set_global_variables from megatron.training.global_vars import set_args, set_global_variables
from megatron.legacy.model import module from megatron.legacy.model import module
from megatron.core import mpu from megatron.core import mpu
from megatron.core.enums import ModelType from megatron.core.enums import ModelType
from megatron.legacy import fused_kernels from megatron.legacy import fused_kernels
except ModuleNotFoundError: except ModuleNotFoundError:
print("Unable to import Megatron, please specify the path to Megatron using --megatron-path. Exiting.") print("Unable to import Megatron, please specify the path to Megatron using --megatron-path. Exiting.")
queue.put("exit") queue.put("exit")
exit(1) exit(1)
# We want all arguments to come from us. # We want all arguments to come from us.
sys.argv = ['script.py', sys.argv = ['script.py',
'--no-masked-softmax-fusion', '--no-masked-softmax-fusion',
'--no-bias-gelu-fusion', '--no-bias-gelu-fusion',
'--no-bias-dropout-fusion', '--no-bias-dropout-fusion',
'--no-async-tensor-model-parallel-allreduce', '--no-async-tensor-model-parallel-allreduce',
'--use-cpu-initialization', '--use-cpu-initialization',
'--micro-batch-size', '1', '--micro-batch-size', '1',
'--no-load-optim', '--no-load-optim',
'--no-load-rng', '--no-load-rng',
'--no-save-optim', '--no-save-optim',
'--no-save-rng', '--no-save-rng',
'--mock-data', # To pass the "blend data checks" in arguments.py '--mock-data', # To pass the "blend data checks" in arguments.py
'--no-initialization', '--no-initialization',
'--load', args.load_dir, '--load', args.load_dir,
'--no-one-logger', '--no-one-logger',
] ]
if args.make_vocab_size_divisible_by is not None: if args.make_vocab_size_divisible_by is not None:
sys.argv.extend(["--make-vocab-size-divisible-by", str(args.make_vocab_size_divisible_by)]) sys.argv.extend(["--make-vocab-size-divisible-by", str(args.make_vocab_size_divisible_by)])
margs = parse_args() margs = parse_args()
margs.tokenizer_model = args.tokenizer_model margs.tokenizer_model = args.tokenizer_model
load_args_from_checkpoint(margs, args.model_size) load_args_from_checkpoint(margs, args.model_size)
if "llama2" in args.model_size: if "llama2" in args.model_size:
margs.tokenizer_type = "Llama2Tokenizer" margs.tokenizer_type = "Llama2Tokenizer"
elif "yi" in args.model_size: elif "yi" in args.model_size:
margs.tokenizer_type = "HuggingFaceTokenizer" margs.tokenizer_type = "HuggingFaceTokenizer"
elif "llama3" in args.model_size: elif "llama3" in args.model_size:
margs.tokenizer_type = "HuggingFaceTokenizer" margs.tokenizer_type = "HuggingFaceTokenizer"
elif "mistral" in args.model_size: elif "mistral" in args.model_size:
margs.tokenizer_type = "HuggingFaceTokenizer" margs.tokenizer_type = "HuggingFaceTokenizer"
elif "qwen2.5" in args.model_size: elif "qwen2.5" in args.model_size:
margs.tokenizer_type = "HuggingFaceTokenizer" margs.tokenizer_type = "HuggingFaceTokenizer"
margs.add_qkv_bias = True margs.add_qkv_bias = True
# Arguments do sanity checks on the world size, but we don't care, # Arguments do sanity checks on the world size, but we don't care,
# so trick it into thinking we are plenty of processes. # so trick it into thinking we are plenty of processes.
margs.world_size = margs.tensor_model_parallel_size * margs.pipeline_model_parallel_size margs.world_size = margs.tensor_model_parallel_size * margs.pipeline_model_parallel_size
margs = validate_args(margs) margs = validate_args(margs)
margs.use_legacy_models = True margs.use_legacy_models = True
margs.transformer_impl = args.loader_transformer_impl margs.transformer_impl = args.loader_transformer_impl
margs.position_embedding_type = "rope" margs.position_embedding_type = "rope"
def check_for_arg(arg_name, default=None): def check_for_arg(arg_name, default=None):
if getattr(margs, arg_name, None) is None: if getattr(margs, arg_name, None) is None:
if default is not None: if default is not None:
setattr(margs, arg_name, default) setattr(margs, arg_name, default)
else: else:
print(f"Checkpoint does not specify the argument {arg_name}. Exiting.") print(f"Checkpoint does not specify the argument {arg_name}. Exiting.")
print(f"Arguments: {margs}") print(f"Arguments: {margs}")
queue.put("exit") queue.put("exit")
exit(1) exit(1)
check_for_arg('tensor_model_parallel_size') check_for_arg('tensor_model_parallel_size')
check_for_arg('pipeline_model_parallel_size') check_for_arg('pipeline_model_parallel_size')
check_for_arg('num_layers') check_for_arg('num_layers')
check_for_arg('hidden_size') check_for_arg('hidden_size')
check_for_arg('seq_length') check_for_arg('seq_length')
check_for_arg('num_attention_heads') check_for_arg('num_attention_heads')
check_for_arg('max_position_embeddings') check_for_arg('max_position_embeddings')
check_for_arg('position_embedding_type') check_for_arg('position_embedding_type')
check_for_arg('iteration') check_for_arg('iteration')
check_for_arg('bert_binary_head') check_for_arg('bert_binary_head')
check_for_arg('disable_bias_linear', False) check_for_arg('disable_bias_linear', False)
check_for_arg('params_dtype') check_for_arg('params_dtype')
check_for_arg('swiglu', False) check_for_arg('swiglu', False)
# Determine how to make our models. # Determine how to make our models.
assert args.model_type == 'GPT', 'Llama-2, Llama-3 and Mistral are GPT models.' assert args.model_type == 'GPT', 'Llama-2, Llama-3 and Mistral are GPT models.'
margs.model_type = ModelType.encoder_or_decoder margs.model_type = ModelType.encoder_or_decoder
margs.params_dtype = torch.bfloat16 if args.bf16 else torch.float16 if args.fp16 else torch.float32 margs.params_dtype = torch.bfloat16 if args.bf16 else torch.float16 if args.fp16 else torch.float32
# Suppress warning about torch.distributed not being initialized. # Suppress warning about torch.distributed not being initialized.
module.MegatronModule.embedding_warning_printed = True module.MegatronModule.embedding_warning_printed = True
set_global_variables(margs, build_tokenizer=False) set_global_variables(margs, build_tokenizer=False)
mpu.set_tensor_model_parallel_world_size(margs.tensor_model_parallel_size) mpu.set_tensor_model_parallel_world_size(margs.tensor_model_parallel_size)
mpu.set_pipeline_model_parallel_world_size(margs.pipeline_model_parallel_size) mpu.set_pipeline_model_parallel_world_size(margs.pipeline_model_parallel_size)
mpu.set_virtual_pipeline_model_parallel_world_size(margs.virtual_pipeline_model_parallel_size) mpu.set_virtual_pipeline_model_parallel_world_size(margs.virtual_pipeline_model_parallel_size)
fused_kernels.load(margs) # fused_kernels.load(margs)
# Short aliases. # Short aliases.
tp_size = margs.tensor_model_parallel_size tp_size = margs.tensor_model_parallel_size
pp_size = margs.pipeline_model_parallel_size pp_size = margs.pipeline_model_parallel_size
vp_size = margs.virtual_pipeline_model_parallel_size vp_size = margs.virtual_pipeline_model_parallel_size
if vp_size is None: if vp_size is None:
vp_size = 1 vp_size = 1
# Metadata. # Metadata.
md = types.SimpleNamespace() md = types.SimpleNamespace()
md.model_type = args.model_type md.model_type = args.model_type
md.num_layers = margs.num_layers md.num_layers = margs.num_layers
md.hidden_size = margs.hidden_size md.hidden_size = margs.hidden_size
md.seq_length = margs.seq_length md.seq_length = margs.seq_length
md.num_attention_heads = margs.num_attention_heads md.num_attention_heads = margs.num_attention_heads
md.max_position_embeddings = margs.max_position_embeddings md.max_position_embeddings = margs.max_position_embeddings
md.tokenizer_type = margs.tokenizer_type md.tokenizer_type = margs.tokenizer_type
md.iteration = margs.iteration md.iteration = margs.iteration
md.params_dtype = margs.params_dtype md.params_dtype = margs.params_dtype
md.bert_binary_head = margs.bert_binary_head md.bert_binary_head = margs.bert_binary_head
md.output_layer = margs.untie_embeddings_and_output_weights md.output_layer = margs.untie_embeddings_and_output_weights
md.position_embedding_type = margs.position_embedding_type md.position_embedding_type = margs.position_embedding_type
md.linear_bias = margs.add_bias_linear md.linear_bias = margs.add_bias_linear
md.qkv_bias = margs.add_qkv_bias md.qkv_bias = margs.add_qkv_bias
md.norm_has_bias = False md.norm_has_bias = False
md.swiglu = margs.swiglu md.swiglu = margs.swiglu
md.previous_tensor_parallel_size = margs.tensor_model_parallel_size md.previous_tensor_parallel_size = margs.tensor_model_parallel_size
md.previous_pipeline_parallel_size = margs.pipeline_model_parallel_size md.previous_pipeline_parallel_size = margs.pipeline_model_parallel_size
md.make_vocab_size_divisible_by = margs.make_vocab_size_divisible_by md.make_vocab_size_divisible_by = margs.make_vocab_size_divisible_by
md.checkpoint_args = margs md.checkpoint_args = margs
md.consumed_train_samples = 0 md.consumed_train_samples = 0
md.consumed_valid_samples = 0 md.consumed_valid_samples = 0
margs.model_size = args.model_size margs.model_size = args.model_size
# Get true (non-padded) vocab size # Get true (non-padded) vocab size
tokenizer = transformers.AutoTokenizer.from_pretrained(margs.tokenizer_model) tokenizer = transformers.AutoTokenizer.from_pretrained(margs.tokenizer_model)
md.true_vocab_size = tokenizer._tokenizer.get_vocab_size(with_added_tokens=True) md.true_vocab_size = tokenizer._tokenizer.get_vocab_size(with_added_tokens=True)
# Get first pipe stage. # Get first pipe stage.
mpu.set_tensor_model_parallel_rank(0) mpu.set_tensor_model_parallel_rank(0)
mpu.set_pipeline_model_parallel_rank(0) mpu.set_pipeline_model_parallel_rank(0)
model = load_checkpoint_to_model(margs) model = load_checkpoint_to_model(margs)
queue.put(md) queue.put(md)
def queue_put(name, msg): def queue_put(name, msg):
print(f"sending {name}") print(f"sending {name}")
msg["name"] = name msg["name"] = name
queue.put(msg) queue.put(msg)
# Send embeddings. # Send embeddings.
message = { message = {
"word embeddings": model.language_model.embedding.word_embeddings.weight.data "word embeddings": model.language_model.embedding.word_embeddings.weight.data
} }
if md.position_embedding_type == 'learned_absolute': if md.position_embedding_type == 'learned_absolute':
message["position embeddings"] = model.language_model.embedding.position_embeddings.weight.data message["position embeddings"] = model.language_model.embedding.position_embeddings.weight.data
else: else:
assert not hasattr(model.language_model.embedding, 'position_embeddings') assert not hasattr(model.language_model.embedding, 'position_embeddings')
queue_put("embeddings", message) queue_put("embeddings", message)
for layer_num in range(margs.num_layers): for layer_num in range(margs.num_layers):
message = {} message = {}
# Get non-parallel tensors from tp_rank 0. # Get non-parallel tensors from tp_rank 0.
layer = model.language_model.encoder.layers[layer_num] layer = model.language_model.encoder.layers[layer_num]
message["input norm weight"] = layer.input_norm.weight.data message["input norm weight"] = layer.input_norm.weight.data
message["post norm weight"] = layer.post_attention_norm.weight.data message["post norm weight"] = layer.post_attention_norm.weight.data
if md.linear_bias: if md.linear_bias:
message["dense bias"] = layer.self_attention.dense.bias.data message["dense bias"] = layer.self_attention.dense.bias.data
message["mlp l1 bias"] = layer.mlp.dense_4h_to_h.bias.data message["mlp l1 bias"] = layer.mlp.dense_4h_to_h.bias.data
# Grab all parallel tensors for this layer. # Grab all parallel tensors for this layer.
qkv_weight = [] qkv_weight = []
qkv_bias = [] qkv_bias = []
dense_weight = [] dense_weight = []
mlp_l0_weight = [] mlp_l0_weight = []
mlp_l0_bias = [] mlp_l0_bias = []
mlp_l1_weight = [] mlp_l1_weight = []
layer = model.language_model.encoder.layers[layer_num] layer = model.language_model.encoder.layers[layer_num]
qkv_weight.append(layer.self_attention.query_key_value.weight.data) qkv_weight.append(layer.self_attention.query_key_value.weight.data)
dense_weight.append(layer.self_attention.dense.weight.data) dense_weight.append(layer.self_attention.dense.weight.data)
mlp_l0_weight.append(layer.mlp.dense_h_to_4h.weight.data) mlp_l0_weight.append(layer.mlp.dense_h_to_4h.weight.data)
mlp_l1_weight.append(layer.mlp.dense_4h_to_h.weight.data) mlp_l1_weight.append(layer.mlp.dense_4h_to_h.weight.data)
if md.qkv_bias: if md.qkv_bias:
qkv_bias.append(layer.self_attention.query_key_value.bias.data) qkv_bias.append(layer.self_attention.query_key_value.bias.data)
if md.linear_bias: if md.linear_bias:
mlp_l0_bias.append(layer.mlp.dense_h_to_4h.bias.data) mlp_l0_bias.append(layer.mlp.dense_h_to_4h.bias.data)
# Handle gated linear units. # Handle gated linear units.
if md.swiglu: if md.swiglu:
# Concat all the first halves ('W's) and all the second halves ('V's). # Concat all the first halves ('W's) and all the second halves ('V's).
for tp_rank in range(tp_size): for tp_rank in range(tp_size):
mlp_l0_weight[tp_rank] = torch.chunk(mlp_l0_weight[tp_rank], 2, dim=0) mlp_l0_weight[tp_rank] = torch.chunk(mlp_l0_weight[tp_rank], 2, dim=0)
message["mlp l0 weight W"] = torch.cat([w[0] for w in mlp_l0_weight], dim=0) message["mlp l0 weight W"] = torch.cat([w[0] for w in mlp_l0_weight], dim=0)
message["mlp l0 weight V"] = torch.cat([w[1] for w in mlp_l0_weight], dim=0) message["mlp l0 weight V"] = torch.cat([w[1] for w in mlp_l0_weight], dim=0)
else: else:
message["mlp l0 weight"] = torch.cat(mlp_l0_weight, dim=0) message["mlp l0 weight"] = torch.cat(mlp_l0_weight, dim=0)
# Simple concat of the rest. # Simple concat of the rest.
message["qkv weight"] = torch.cat(qkv_weight, dim=0) message["qkv weight"] = torch.cat(qkv_weight, dim=0)
message["dense weight"] = torch.cat(dense_weight, dim=1) message["dense weight"] = torch.cat(dense_weight, dim=1)
message["mlp l1 weight"] = torch.cat(mlp_l1_weight, dim=1) message["mlp l1 weight"] = torch.cat(mlp_l1_weight, dim=1)
if md.qkv_bias: if md.qkv_bias:
message["qkv bias"] = torch.cat(qkv_bias, dim=0) message["qkv bias"] = torch.cat(qkv_bias, dim=0)
if md.linear_bias: if md.linear_bias:
if md.swiglu: if md.swiglu:
for tp_rank in range(tp_size): for tp_rank in range(tp_size):
mlp_l0_bias[tp_rank] = torch.chunk(mlp_l0_bias[tp_rank], 2, dim=0) mlp_l0_bias[tp_rank] = torch.chunk(mlp_l0_bias[tp_rank], 2, dim=0)
message["mlp l0 bias W"] = torch.cat([b[0] for b in mlp_l0_bias],dim=0) message["mlp l0 bias W"] = torch.cat([b[0] for b in mlp_l0_bias],dim=0)
message["mlp l0 bias V"] = torch.cat([b[1] for b in mlp_l0_bias],dim=0) message["mlp l0 bias V"] = torch.cat([b[1] for b in mlp_l0_bias],dim=0)
else: else:
message["mlp l0 bias"] = torch.cat(mlp_l0_bias, dim=0) message["mlp l0 bias"] = torch.cat(mlp_l0_bias, dim=0)
queue_put(f"transformer layer {layer_num}", message) queue_put(f"transformer layer {layer_num}", message)
# Send final norm from tp_rank 0. # Send final norm from tp_rank 0.
message = { message = {
"weight": model.language_model.encoder.final_norm.weight.data, "weight": model.language_model.encoder.final_norm.weight.data,
} }
queue_put("final norm", message) queue_put("final norm", message)
if md.output_layer: if md.output_layer:
message = { message = {
"weight": model.language_model.output_layer.weight.data "weight": model.language_model.output_layer.weight.data
} }
queue_put("output layer", message) queue_put("output layer", message)
queue.put("done") queue.put("done")
if args.checkpoint_type == "meta": if args.checkpoint_type == "meta":
shutil.rmtree(os.path.join(args.load_dir)) shutil.rmtree(os.path.join(args.load_dir))
def load_checkpoint(queue, args): def load_checkpoint(queue, args):
try: try:
_load_checkpoint(queue, args) _load_checkpoint(queue, args)
except Exception: except Exception:
queue.put("exit") queue.put("exit")
raise raise
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment