import os import shutil import torch import json import logging from torch import distributed as dist from safetensors import safe_open from transformers.modeling_utils import SAFE_WEIGHTS_INDEX_NAME from huggingface_hub.constants import SAFETENSORS_INDEX_FILE from megatron.training.checkpointing import ( save_checkpoint, get_checkpoint_tracker_filename, get_checkpoint_name ) from general.synchronizer import BaseSynchronizer, ParamType class HF2MGSynchronizer(BaseSynchronizer): def __init__(self, load_dir, model_provider_func=None): super().__init__(load_dir, model_provider_func) self._single_file = False p = os.path.join(self.load_dir, SAFE_WEIGHTS_INDEX_NAME) if not os.path.exists(p): p = os.path.join(self.load_dir, SAFETENSORS_INDEX_FILE) if os.path.exists(p): with open(p, 'r') as f: data = json.load(f)['weight_map'] self._key_to_file = {k: os.path.join(self.load_dir, v) for k, v in data.items()} elif os.path.exists(os.path.join(self.load_dir, 'model.safetensors')): self._single_file = True self._key_to_file = dict() else: raise FileNotFoundError() if self.debug: if not self.dryrun: for p in self._mgmodel.parameters(): p.data.fill_(torch.nan) # NOTE: Fill non-persistent/persistent buffer with NaN for b in self._mgmodel.buffers(): b.data.fill_(torch.nan) self._visit = torch.zeros([self.hf_size], dtype=torch.int, device=self.device) def load_tensor(self, dummy_tensor): def _get_filename_from_key(key): if self._single_file: return os.path.join(self.load_dir, 'model.safetensors') if key in self._key_to_file: return self._key_to_file[key] raise KeyError(f'{key} not found in index file') if dummy_tensor not in self._hf_params_to_key: raise ValueError() key = self._hf_params_to_key[dummy_tensor] if self.debug: self._visit[self._hf_params_key_to_id[key]] = True if not self.args.untie_embeddings_and_output_weights and key == 'lm_head.weight': key = 'model.embed_tokens.weight' file = _get_filename_from_key(key) with safe_open(file, framework="pt", device=str(self.device)) as f: return f.get_tensor(key) def _copy_impl( self, src_tensor, dst_tensor, param_type: ParamType=ParamType.UNIQUE ): tp_rank, tp_size = self.tp_rank, self.tp_size if param_type in [ParamType.MOE_COLUMN, ParamType.MOE_ROW, ParamType.MOE_GATE_UP]: tp_rank, tp_size = self.etp_rank, self.etp_size split_mapping = { ParamType.UNIQUE: lambda x: self.load_tensor(x), ParamType.COLUMN: lambda x: torch.chunk(self.load_tensor(x), tp_size, dim=0)[tp_rank], ParamType.ROW: lambda x: torch.chunk(self.load_tensor(x), tp_size, dim=1)[tp_rank], # the data of following type is loaded by caller ParamType.GATE_UP: lambda x: torch.chunk(x, tp_size, dim=1)[tp_rank].flatten(0, 1), ParamType.QKV_W: lambda x: torch.chunk(x, tp_size, dim=0)[tp_rank].flatten(0, 1), ParamType.QKV_B: lambda x: torch.chunk(x, tp_size, dim=0)[tp_rank].flatten(), ParamType.MOE_COLUMN: lambda x: torch.chunk(self.load_tensor(x), tp_size, dim=0)[tp_rank], ParamType.MOE_ROW: lambda x: torch.chunk(self.load_tensor(x), tp_size, dim=1)[tp_rank], # the data of following type is loaded by caller ParamType.MOE_GATE_UP: lambda x: torch.chunk(x, tp_size, dim=1)[tp_rank].flatten(0, 1), } if self.dryrun: return dst_tensor.data.copy_(dst_tensor.clone()) dst_tensor.data.copy_(split_mapping[param_type](src_tensor)) def set_preprocess_state(self): '''Set embedding params.''' self.copy( self._hfmodel.model.embed_tokens.weight, self._mgmodel.embedding.word_embeddings.weight, param_type=ParamType.COLUMN ) def set_postprocess_state(self): '''Set output layer & norm params.''' self.copy( self._hfmodel.model.norm.weight, self._mgmodel.decoder.final_layernorm.weight, ) if self._mgmodel.share_embeddings_and_output_weights: output_layer_weight = self._mgmodel.shared_embedding_or_output_weight() else: output_layer_weight = self._mgmodel.output_layer.weight self.copy( self._hfmodel.lm_head.weight, output_layer_weight, param_type=ParamType.COLUMN ) def set_mla_selfattn_state(self, attn, hf_attn): # NOTE: MLA qkv_bias always False if self.args.q_lora_rank is None: self.copy(hf_attn.q_proj.weight, attn.linear_q_proj.weight, param_type=ParamType.COLUMN) else: self.copy(hf_attn.q_a_proj.weight, attn.linear_q_down_proj.weight, param_type=ParamType.COLUMN) self.copy(hf_attn.q_b_proj.weight, attn.linear_q_up_proj.weight, param_type=ParamType.COLUMN) if self.args.qk_layernorm: self.copy( hf_attn.q_a_layernorm.weight, attn.linear_q_up_proj.layer_norm_weight ) self.copy(hf_attn.kv_a_proj_with_mqa.weight, attn.linear_kv_down_proj.weight, param_type=ParamType.COLUMN) self.copy(hf_attn.kv_b_proj.weight, attn.linear_kv_up_proj.weight, param_type=ParamType.COLUMN) if self.args.qk_layernorm: self.copy( hf_attn.kv_a_layernorm.weight, attn.linear_kv_up_proj.layer_norm_weight ) self.copy( hf_attn.o_proj.weight, attn.linear_proj.weight, param_type=ParamType.ROW ) def set_selfattn_state(self, attn, hf_attn): '''Set self-attention params.''' # Reshape loaded weights. num_heads = self.args.num_attention_heads num_query_groups = (self.args.num_query_groups if self.args.group_query_attention else self.args.num_attention_heads) num_querys_per_group = num_heads // num_query_groups dim = self.args.kv_channels assert num_heads % num_querys_per_group == 0 # copy qk norm if indeed. if self.args.qk_layernorm: self.copy(hf_attn.q_norm.weight, attn.q_layernorm.weight) self.copy(hf_attn.k_norm.weight, attn.k_layernorm.weight) # Copy weights (re-order dimensions for Megatron). if self.dryrun: attn_proj_weight = attn.linear_qkv.weight else: attn_proj_weight = torch.cat([ self.load_tensor(hf_attn.q_proj.weight).reshape((num_query_groups, num_querys_per_group*dim, -1)), self.load_tensor(hf_attn.k_proj.weight).reshape((num_query_groups, dim, -1)), self.load_tensor(hf_attn.v_proj.weight).reshape((num_query_groups, dim, -1)), ], dim=1) self.copy( attn_proj_weight, attn.linear_qkv.weight, param_type=ParamType.QKV_W, ) self.copy( hf_attn.o_proj.weight, attn.linear_proj.weight, param_type=ParamType.ROW ) # Copy bias if self.args.add_qkv_bias: if self.dryrun: attn_proj_bias = attn.linear_qkv.bias else: attn_proj_bias = torch.cat([ self.load_tensor(hf_attn.q_proj.bias).reshape((num_query_groups, num_querys_per_group*dim, -1)), self.load_tensor(hf_attn.k_proj.bias).reshape((num_query_groups, dim, -1)), self.load_tensor(hf_attn.v_proj.bias).reshape((num_query_groups, dim, -1)), ], dim=1) self.copy( attn_proj_bias, attn.linear_qkv.bias, param_type=ParamType.QKV_B, ) def set_mlp_state(self, mlp, hf_mlp, expert_id=''): ''' Set MLP params. The mlp (mcore MLP) should have attributes `linear_fc1` and `linear_fc2`. Currently only Gated Linear is supported. ''' if self.dryrun: gate_up_proj_weight = mlp.linear_fc1.weight else: gate_up_proj_weight = torch.stack([ self.load_tensor(hf_mlp.gate_proj.weight), self.load_tensor(hf_mlp.up_proj.weight) ]) linear_fc1_weight = getattr(mlp.linear_fc1, f'weight{expert_id}') linear_fc2_weight = getattr(mlp.linear_fc2, f'weight{expert_id}') self.copy( gate_up_proj_weight, linear_fc1_weight, param_type=ParamType.GATE_UP if expert_id == '' else ParamType.MOE_GATE_UP ) self.copy( hf_mlp.down_proj.weight, linear_fc2_weight, param_type=ParamType.ROW if expert_id == '' else ParamType.MOE_ROW ) def set_sequential_mlp_state(self, experts, hf_experts): '''Set MOE MLP params.''' experts = experts.local_experts for mg_expert_id, hf_expert_id in self._build_expert_parallel_mapping().items(): self.set_mlp_state(experts[mg_expert_id], hf_experts[hf_expert_id]) def set_group_mlp_state(self, experts, hf_experts): for mg_expert_id, hf_expert_id in self._build_expert_parallel_mapping().items(): self.set_mlp_state(experts, hf_experts[hf_expert_id], expert_id=mg_expert_id) def set_moe_layer_state(self, moe, hf_moe): # router self.copy(hf_moe.gate.weight, moe.router.weight) if moe.router.enable_expert_bias: self.copy(hf_moe.gate.e_score_correction_bias, moe.router.expert_bias) # experts if self.args.moe_grouped_gemm: # group gemm if self.args.moe_use_legacy_grouped_gemm: # weight1 and weight2, not impl raise NotImplementedError("Currently only TE GroupGEMM is implemented.") self.set_group_mlp_state(moe.experts, hf_moe.experts) else: # sequential self.set_sequential_mlp_state(moe.experts, hf_moe.experts) # shared experts if moe.shared_experts is not None: if moe.shared_experts.use_shared_expert_gate: self.copy(hf_moe.shared_expert_gate.weight, moe.shared_experts.gate_weight) self.set_mlp_state(moe.shared_experts, hf_moe.shared_experts) def set_layer_state(self, layer, hf_layer): '''Set transformer layer params.''' if self.args.multi_latent_attention: self.set_mla_selfattn_state(layer.self_attention, hf_layer.self_attn) self.copy(hf_layer.input_layernorm.weight, layer.input_layernorm.weight) else: self.set_selfattn_state(layer.self_attention, hf_layer.self_attn) self.copy(hf_layer.input_layernorm.weight, layer.self_attention.linear_qkv.layer_norm_weight) if hasattr(layer.mlp, 'router'): self.set_moe_layer_state(layer.mlp, hf_layer.mlp) self.copy(hf_layer.post_attention_layernorm.weight, layer.pre_mlp_layernorm.weight) else: self.set_mlp_state(layer.mlp, hf_layer.mlp) self.copy(hf_layer.post_attention_layernorm.weight, layer.mlp.linear_fc1.layer_norm_weight) def check_and_save(self, output_dir): if self.debug: if not self.dryrun: for n, p in self._mgmodel.state_dict().items(): if isinstance(p, torch.Tensor) and p.isnan().any(): raise SystemError(f'NaN Parameters Detected on key {n}') from torch.distributed import all_reduce all_reduce(self._visit) unvisit_param_ids = (self._visit == 0).nonzero()[:, 0].cpu().numpy().tolist() unvisit_keys = [] for param_id in unvisit_param_ids: unvisit_keys.append(self._id_to_hf_params_key[param_id]) if len(unvisit_keys) > 0: logging.warning(f"Never visit the following huggingface weights in the conversion: {unvisit_keys}") self.args.save = output_dir if not self.dryrun: save_checkpoint( getattr(self.args, 'iteration', 1), [self._mgmodel], None, None, 0, pipeline_rank=self.pp_rank, pipeline_parallel=self.pp_size > 1, expert_rank=self.ep_rank, expert_parallel=self.ep_size > 1, tensor_rank=self.tp_rank ) dist.barrier() if self.rank == 0: # NOTE: The `save_checkpoint` API can only save a checkpoint in release=False, # reset the metadata. (Otherwise user may find their training starts at step 2) tracker_filename = get_checkpoint_tracker_filename(self.args.save) with open(tracker_filename, 'w') as f: f.write('release') source_dir = get_checkpoint_name(self.args.save, 1, False, return_base_dir=True) target_dir = get_checkpoint_name(self.args.save, -1, True, return_base_dir=True) shutil.move(source_dir, target_dir)