in toolkits/model_checkpoints_convertor/qwen/hf2mcore_qwen2_moe.py [0:0]
def load_megatron_model(args):
os.makedirs(args.save, exist_ok=True)
os.system("cp -rf " + args.hf_ckpt_path + "/*config.json " + args.save)
os.system("cp -rf " + args.hf_ckpt_path + "/tokenizer* " + args.save)
# os.system("cp -rf " + args.hf_ckpt_path + "/*.py " + args.save)
os.system("cp -rf " + args.hf_ckpt_path + "/merges.txt " + args.save)
os.system("cp -rf " + args.hf_ckpt_path + "/*config.json " + args.load)
os.system("cp -rf " + args.hf_ckpt_path + "/tokenizer* " + args.load)
# os.system("cp -rf " + args.hf_ckpt_path + "/*.py " + args.load)
os.system("cp -rf " + args.hf_ckpt_path + "/vocab.json " + args.load)
os.system("cp -rf " + args.hf_ckpt_path + "/merges.txt " + args.load)
model = model_provider()
args.tensor_model_parallel_size = args.target_tensor_model_parallel_size
args.pipeline_model_parallel_size = args.target_pipeline_model_parallel_size
args.expert_tensor_parallel_size = args.target_expert_tensor_parallel_size
if args.num_experts is not None:
args.expert_model_parallel_size = args.target_expert_model_parallel_size
if args.tensor_model_parallel_size > 1:
args.sequence_parallel = True
model_path = args.load
tracker_filename = get_checkpoint_tracker_filename(model_path)
iteration, release = read_metadata(tracker_filename)
head_dim = args.hidden_size // args.num_attention_heads if args.kv_channels is None else args.kv_channels
group_per_split = args.num_query_groups // args.tensor_model_parallel_size
if args.num_experts is not None:
num_local_experts = args.num_experts // args.expert_model_parallel_size
state_dict = {}
mid_state = defaultdict(list)
if (
args.tensor_model_parallel_size >= 1
and args.pipeline_model_parallel_size >= 1
and args.expert_model_parallel_size >= 1
and args.num_experts % args.expert_model_parallel_size == 0
and args.expert_tensor_parallel_size == 1
):
if args.target_decoder_first_pipeline_num_layers is not None:
remained_layers = args.num_layers - args.target_decoder_first_pipeline_num_layers
remained_stages = args.pipeline_model_parallel_size - 1
assert remained_layers % remained_stages == 0
pp_layers_per_stage = [args.target_decoder_first_pipeline_num_layers] +([remained_layers // remained_stages] * remained_stages)
else:
pp_layers_per_stage = [args.num_layers // args.pipeline_model_parallel_size] * args.pipeline_model_parallel_size
layers_to_copy = {}
for tp_rank in range(args.tensor_model_parallel_size):
for ep_rank in range(tp_rank, args.expert_model_parallel_size, args.tensor_model_parallel_size):
for pp_rank in range(args.pipeline_model_parallel_size):
layer_offset = sum(pp_layers_per_stage[:pp_rank])
for layer in range(pp_layers_per_stage[pp_rank]):
pp_layer_id = layer + layer_offset
layers_to_copy[(pp_rank, layer)] = pp_layer_id
if args.expert_model_parallel_size > 1:
checkpoint_name = get_checkpoint_name(model_path, iteration, release, True, tp_rank, pp_rank, True,
ep_rank)
elif args.expert_model_parallel_size == 1:
checkpoint_name = get_checkpoint_name(model_path, iteration, release, True, tp_rank, pp_rank,
False)
print(f'load {checkpoint_name}')
split_state = torch.load(checkpoint_name, map_location="cpu", weights_only=False)['model']
for k, v in split_state.items():
if "_extra_state" in k:
continue
try:
if 'experts' in k and "shared_experts" not in k:
pattern = r'weight(\d+)'
local_expert_rank = int(re.findall(pattern, k)[0])
expert_rank = local_expert_rank + num_local_experts * ep_rank
k = k.replace(f'weight{local_expert_rank}', f'weight{expert_rank}')
pattern = re.compile(r'\d+')
res = pattern.findall(k)
tgt = re.sub(r"decoder.layers.\d+", "decoder.layers." + str(layers_to_copy[(pp_rank, int(res[0]))]), k)
if 'linear_proj' in k or 'shared_experts.linear_fc1' in k or 'shared_experts.linear_fc2' in k or \
"linear_qkv" in k:
if ep_rank == tp_rank:
mid_state[tgt].append(v)
else:
mid_state[tgt].append(v)
except:
if "word_embeddings" in k:
if ep_rank == tp_rank and pp_rank == 0:
mid_state[k].append(v)
elif "output_layer" in k or "final_layernorm" in k:
if ep_rank == tp_rank and pp_rank == args.pipeline_model_parallel_size - 1:
mid_state[k].append(v)
else:
raise ValueError(f"{k} is missing!! ")
for k, v in mid_state.items():
if 'extra_state' in k:
continue
elif not isinstance(v[0], torch.Tensor) or 'router' in k or 'gate' in k:
target_v = v[0]
elif 'word_embeddings' in k or 'output_layer' in k:
target_v = torch.cat(v, dim=0)
elif 'linear_proj' in k:
target_v = torch.cat(v, dim=1)
elif 'linear_qkv.weight' in k:
viewed = [x.view(group_per_split, -1, head_dim, args.hidden_size) for x in v]
target_v = torch.cat(viewed, dim=0).view(-1, args.hidden_size)
elif 'linear_qkv.bias' in k:
viewed = [x.view(group_per_split, -1) for x in v]
target_v = torch.cat(viewed, dim=0).view(-1)
elif "experts.linear_fc2" in k and "shared_experts" not in k:
target_v = v[0]
elif 'experts.linear_fc1' in k and "shared_experts" not in k:
target_v = v[0]
elif 'shared_experts.linear_fc1' in k:
viewed = [x.view(2, -1, args.hidden_size) for x in v]
target_v = torch.cat(viewed, dim=1).view(-1, args.hidden_size)
elif "shared_experts.linear_fc2" in k:
target_v = torch.cat(v, dim=1)
elif "shared_experts.gate_weight" in k or 'layer_norm_weight' in k or 'pre_mlp_layernorm' in k or 'final_layernorm' in k or 'q_layernorm' in k or 'k_layernorm' in k:
target_v = v[0]
else:
raise ValueError(f"{k} is missing!")
state_dict[k] = target_v
elif (
args.tensor_model_parallel_size >= 1
and args.pipeline_model_parallel_size >= 1
and args.expert_model_parallel_size >= 1
and args.num_experts % args.expert_model_parallel_size == 0
and args.expert_tensor_parallel_size > 1
):
if args.target_decoder_first_pipeline_num_layers is not None:
remained_layers = args.num_layers - args.target_decoder_first_pipeline_num_layers
remained_stages = args.pipeline_model_parallel_size - 1
assert remained_layers % remained_stages == 0
pp_layers_per_stage = [args.target_decoder_first_pipeline_num_layers] +([remained_layers // remained_stages] * remained_stages)
else:
pp_layers_per_stage = [args.num_layers // args.pipeline_model_parallel_size] * args.pipeline_model_parallel_size
layers_to_copy = {}
for tp_rank in range(args.tensor_model_parallel_size):
for ep_rank in range(args.expert_model_parallel_size):
for pp_rank in range(args.pipeline_model_parallel_size):
layer_offset = sum(pp_layers_per_stage[:pp_rank])
for layer in range(pp_layers_per_stage[pp_rank]):
pp_layer_id = layer + layer_offset
layers_to_copy[(pp_rank, layer)] = pp_layer_id
if args.expert_model_parallel_size > 1:
checkpoint_name = get_checkpoint_name(model_path, iteration, release, True, tp_rank, pp_rank, True,
ep_rank)
elif args.expert_model_parallel_size == 1:
checkpoint_name = get_checkpoint_name(model_path, iteration, release, True, tp_rank, pp_rank,
False)
print(f'load {checkpoint_name}')
split_state = torch.load(checkpoint_name, map_location="cpu", weights_only=False)['model']
for k, v in split_state.items():
if "_extra_state" in k:
continue
try:
if 'experts' in k and "shared_experts" not in k:
pattern = r'weight(\d+)'
local_expert_rank = int(re.findall(pattern, k)[0])
expert_rank = local_expert_rank + num_local_experts * ep_rank
k = k.replace(f'weight{local_expert_rank}', f'weight{expert_rank}')
pattern = re.compile(r'\d+')
res = pattern.findall(k)
tgt = re.sub(r"decoder.layers.\d+", "decoder.layers." + str(layers_to_copy[(pp_rank, int(res[0]))]), k)
if 'linear_proj' in k or 'shared_experts.linear_fc1' in k or 'shared_experts.linear_fc2' in k or \
"linear_qkv" in k:
if ep_rank == 0:
mid_state[tgt].append(v)
else:
mid_state[tgt].append(v)
except:
if "word_embeddings" in k:
if ep_rank == 0 and pp_rank == 0:
mid_state[k].append(v)
elif "output_layer" in k or "final_layernorm" in k:
if ep_rank == 0 and pp_rank == args.pipeline_model_parallel_size - 1:
mid_state[k].append(v)
else:
raise ValueError(f"{k} is missing!! ")
for k, v in mid_state.items():
if 'extra_state' in k:
continue
elif not isinstance(v[0], torch.Tensor) or 'router' in k or 'gate' in k:
target_v = v[0]
elif 'word_embeddings' in k or 'output_layer' in k:
target_v = torch.cat(v, dim=0)
elif 'linear_proj' in k:
target_v = torch.cat(v, dim=1)
elif 'linear_qkv.weight' in k:
viewed = [x.view(group_per_split, -1, head_dim, args.hidden_size) for x in v]
target_v = torch.cat(viewed, dim=0).view(-1, args.hidden_size)
elif 'linear_qkv.bias' in k:
viewed = [x.view(group_per_split, -1) for x in v]
target_v = torch.cat(viewed, dim=0).view(-1)
elif 'linear_fc1' in k:
viewed = [x.view(2, -1, args.hidden_size) for x in v]
target_v = torch.cat(viewed, dim=1).view(-1, args.hidden_size)
elif "linear_fc2" in k:
target_v = torch.cat(v, dim=1)
elif "shared_experts.gate_weight" in k or 'layer_norm_weight' in k or 'pre_mlp_layernorm' in k or 'final_layernorm' in k or 'q_layernorm' in k or 'k_layernorm' in k:
target_v = v[0]
else:
raise ValueError(f"{k} is missing!")
state_dict[k] = target_v
else:
raise ValueError('not support yet')
model.load_state_dict(state_dict, strict=False)
return model