in toolkits/model_checkpoints_convertor/qwen/hf2megablocks_qwen1.5.py [0:0]
def load_megatron_model(args, model):
model_path = args.load_path
tracker_filename = get_checkpoint_tracker_filename(model_path)
iteration, release = read_metadata(tracker_filename)
head_dim = args.hidden_size // args.num_attention_heads
group_per_split = args.num_query_groups // args.target_tensor_model_parallel_size
num_local_experts = args.moe_num_experts // args.target_expert_model_parallel_size if args.moe_num_experts else 0
state_dict = {}
mid_state = defaultdict(list)
if (
args.target_tensor_model_parallel_size == 1
and args.target_pipeline_model_parallel_size == 1
and args.target_expert_model_parallel_size == 1
):
checkpoint_name = get_checkpoint_names(model_path, iteration, release, None, None, None, None, None)
state_dict = torch.load(checkpoint_name)['model']
elif (
args.target_tensor_model_parallel_size == 1
and args.target_pipeline_model_parallel_size == 1
and args.moe_num_experts
and args.moe_num_experts % args.target_expert_model_parallel_size == 0
):
for ep_rank in range(args.target_expert_model_parallel_size):
checkpoint_name = get_checkpoint_names(model_path, iteration, release, None, None, None, True, ep_rank)
print(f'load {checkpoint_name}')
split_state = torch.load(checkpoint_name, map_location="cpu")['model']
for k, v in split_state.items():
if 'local_experts' in k:
expert_local_rank = name_to_expert_rank(k)
expert_rank = expert_local_rank + num_local_experts * ep_rank
k = k.replace(f'local_experts.{expert_local_rank}', f'local_experts.{expert_rank}')
state_dict[k] = v
elif (
args.target_tensor_model_parallel_size > 1
and args.target_pipeline_model_parallel_size == 1
and args.moe_num_experts is None
):
for tp_rank in range(args.target_tensor_model_parallel_size):
checkpoint_name = get_checkpoint_names(model_path, iteration, release, None, tp_rank, None, None, None)
print(f'load {checkpoint_name}')
split_state = torch.load(checkpoint_name, map_location="cpu")['model']
for k, v in split_state.items():
mid_state[k].append(v)
for k, v in mid_state.items():
if not isinstance(v[0], torch.Tensor) or 'norm' in k:
target_v = v[0]
elif 'embedding' in k or 'output_layer' in k:
target_v = torch.cat(v, dim=0)
elif 'linear_proj' in k or 'linear_fc2' 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)
else:
raise ValueError
state_dict[k] = target_v
elif (
args.target_tensor_model_parallel_size > 1
and args.target_pipeline_model_parallel_size == 1
and args.moe_num_experts
and args.moe_num_experts % args.target_expert_model_parallel_size == 0
):
for tp_rank in range(args.target_tensor_model_parallel_size):
for ep_rank in range(args.target_expert_model_parallel_size):
checkpoint_name = get_checkpoint_names(model_path, iteration, release, None, tp_rank, None, True, ep_rank)
print(f'load {checkpoint_name}')
split_state = torch.load(checkpoint_name, map_location="cpu")['model']
for k, v in split_state.items():
if 'local_experts' in k and 'norm' not in k:
local_expert_rank = name_to_expert_rank(k)
expert_rank = local_expert_rank + num_local_experts * ep_rank
k = k.replace(f'local_experts.{local_expert_rank}', f'local_experts.{expert_rank}')
mid_state[k].append(v)
elif ep_rank == 0:
mid_state[k].append(v)
for k, v in mid_state.items():
if not isinstance(v[0], torch.Tensor) or 'norm' in k or 'router' in k:
target_v = v[0]
elif 'embedding' in k or 'output_layer' in k:
target_v = torch.cat(v, dim=0)
elif 'linear_proj' in k or 'linear_fc2' 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)
else:
print('passed', k)
state_dict[k] = target_v
else:
raise ValueError('not support yet')
model.load_state_dict(state_dict)
return model