import os import torch import json import logging from typing import * from collections import defaultdict from functools import partial from torch import distributed as dist from safetensors.torch import save_file as safe_save_file from huggingface_hub.serialization import split_torch_state_dict_into_shards from megatron.training.checkpointing import load_checkpoint from general.synchronizer import BaseSynchronizer, ParamType class ParamMergeError(ValueError): ... class MG2HFSynchronizer(BaseSynchronizer): def __init__(self, load_dir, model_provider_func=None): super().__init__(load_dir, model_provider_func) if not self.dryrun: load_checkpoint( [self._mgmodel], None, None, checkpointing_context=None, skip_load_to_model_and_opt=False ) self.num_savers = self.args.num_hf_saver self.max_shard_size = self.args.max_shard_size # NOTE: mapping unique global id (0 ~ n-1) to local data self._local_params = dict() # NOTE: mapping global id to global metadata self._tensor_shape = dict() self._tensor_dtype = dict() # sharded shape (rather than hf param shape) # mapping rank to (tp, pp, etp, ep, dp, edp) self._rank_mapping = torch.zeros([self.world_size, 6], dtype=torch.int, device=self.device) self._rank_mapping[self.rank] = torch.Tensor([self.tp_rank, self.pp_rank, self.etp_rank, self.ep_rank, self.dp_rank, self.edp_rank]).to(self.device) dist.all_gather_into_tensor(self._rank_mapping, self._rank_mapping[self.rank]) # define the merge function type for each param self._merge_type: torch.Tensor = torch.zeros([self.hf_size], dtype=torch.int, device=self.device) self._has_param: torch.Tensor = None # self._has_param[param_id].nonzero() ==> ranks that have this param def _copy_impl(self, src_tensor, dst_tensor, param_type: ParamType=ParamType.UNIQUE): param_id = self._hf_params_to_id[dst_tensor] if param_type in [ParamType.MOE_COLUMN, ParamType.MOE_ROW, ParamType.MOE_GATE_UP]: # NOTE: only register on edp_rank 0 if self.edp_rank != 0: return elif param_type == ParamType.UNIQUE: if self.dp_rank != 0 or self.edp_rank != 0: return elif self.dp_rank != 0: return self._local_params[param_id] = src_tensor self._merge_type[param_id] = param_type.value def set_preprocess_state(self): '''Set embedding params.''' self.copy( self._mgmodel.embedding.word_embeddings.weight, self._hfmodel.model.embed_tokens.weight, param_type=ParamType.COLUMN ) def set_postprocess_state(self): '''Set output layer & norm params.''' self.copy(self._mgmodel.decoder.final_layernorm.weight, self._hfmodel.model.norm.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(output_layer_weight, self._hfmodel.lm_head.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(attn.linear_q_proj.weight, hf_attn.q_proj.weight, param_type=ParamType.COLUMN) else: self.copy(attn.linear_q_down_proj.weight, hf_attn.q_a_proj.weight, param_type=ParamType.COLUMN) self.copy(attn.linear_q_up_proj.weight, hf_attn.q_b_proj.weight, param_type=ParamType.COLUMN) if self.args.qk_layernorm: self.copy( attn.linear_q_up_proj.layer_norm_weight, hf_attn.q_a_layernorm.weight ) self.copy(attn.linear_kv_down_proj.weight, hf_attn.kv_a_proj_with_mqa.weight, param_type=ParamType.COLUMN) self.copy(attn.linear_kv_up_proj.weight, hf_attn.kv_b_proj.weight, param_type=ParamType.COLUMN) if self.args.qk_layernorm: self.copy( attn.linear_kv_up_proj.layer_norm_weight, hf_attn.kv_a_layernorm.weight ) self.copy( attn.linear_proj.weight, hf_attn.o_proj.weight, param_type=ParamType.ROW ) def set_selfattn_state(self, attn, hf_attn): '''Set self-attention params.''' # Reshape loaded weights. tp = self.tp_size 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(attn.q_layernorm.weight, hf_attn.q_norm.weight) self.copy(attn.k_layernorm.weight, hf_attn.k_norm.weight) # Copy weights (re-order dimensions for Megatron). attn_proj_weight = attn.linear_qkv.weight.reshape( (num_query_groups // tp, (2 + num_querys_per_group)*dim, -1) ) ( q_proj_weight, k_proj_weight, v_proj_weight ) = torch.split(attn_proj_weight, [num_querys_per_group*dim, dim, dim], dim=1) self.copy(q_proj_weight, hf_attn.q_proj.weight, param_type=ParamType.QKV_W) self.copy(k_proj_weight, hf_attn.k_proj.weight, param_type=ParamType.QKV_W) self.copy(v_proj_weight, hf_attn.v_proj.weight, param_type=ParamType.QKV_W) self.copy( attn.linear_proj.weight, hf_attn.o_proj.weight, param_type=ParamType.ROW ) # Copy bias if self.args.add_qkv_bias: attn_proj_bias = attn.linear_qkv.bias.reshape( (num_query_groups // tp, (2 + num_querys_per_group)*dim, -1) ) q_proj_bias, k_proj_bias, v_proj_bias = torch.split( attn_proj_bias, [num_querys_per_group*dim, dim, dim], dim=1 ) self.copy(q_proj_bias, hf_attn.q_proj.bias, param_type=ParamType.QKV_B) self.copy(k_proj_bias, hf_attn.k_proj.bias, param_type=ParamType.QKV_B) self.copy(v_proj_bias, hf_attn.v_proj.bias, param_type=ParamType.QKV_B) def set_mlp_state(self, mlp, hf_mlp): '''Set MLP params.''' hidden_size = mlp.linear_fc1.weight.shape[-1] gate_proj_weight, up_proj_weight = mlp.linear_fc1.weight.reshape(2, -1, hidden_size) self.copy( gate_proj_weight, hf_mlp.gate_proj.weight, param_type=ParamType.COLUMN ) self.copy( up_proj_weight, hf_mlp.up_proj.weight, param_type=ParamType.COLUMN ) self.copy( mlp.linear_fc2.weight, hf_mlp.down_proj.weight, param_type=ParamType.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(): hidden_size = experts[mg_expert_id].linear_fc1.weight.shape[-1] ( gate_proj_weight, up_proj_weight ) = experts[mg_expert_id].linear_fc1.weight.reshape(2, -1, hidden_size) self.copy( gate_proj_weight, hf_experts[hf_expert_id].gate_proj.weight, param_type=ParamType.MOE_COLUMN ) self.copy( up_proj_weight, hf_experts[hf_expert_id].up_proj.weight, param_type=ParamType.MOE_COLUMN ) self.copy( experts[mg_expert_id].linear_fc2.weight, hf_experts[hf_expert_id].down_proj.weight, param_type=ParamType.MOE_ROW ) def set_group_mlp_state(self, experts, hf_experts): for mg_expert_id, hf_expert_id in self._build_expert_parallel_mapping().items(): hidden_size = getattr(experts.linear_fc1, f'weight{mg_expert_id}').shape[-1] ( gate_proj_weight, up_proj_weight ) = getattr(experts.linear_fc1, f'weight{mg_expert_id}').reshape(2, -1, hidden_size) self.copy( gate_proj_weight, hf_experts[hf_expert_id].gate_proj.weight, param_type=ParamType.MOE_COLUMN ) self.copy( up_proj_weight, hf_experts[hf_expert_id].up_proj.weight, param_type=ParamType.MOE_COLUMN ) self.copy( getattr(experts.linear_fc2, f'weight{mg_expert_id}'), hf_experts[hf_expert_id].down_proj.weight, param_type=ParamType.MOE_ROW ) def set_moe_layer_state(self, moe, hf_moe): # router self.copy(moe.router.weight, hf_moe.gate.weight) if moe.router.enable_expert_bias: self.copy(moe.router.expert_bias, hf_moe.gate.e_score_correction_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(moe.shared_experts.gate_weight, hf_moe.shared_expert_gate.weight) hidden_size = moe.shared_experts.linear_fc1.weight.shape[-1] gate_proj_weight, up_proj_weight = moe.shared_experts.linear_fc1.weight.reshape(2, -1, hidden_size) self.copy( gate_proj_weight, hf_moe.shared_experts.gate_proj.weight, param_type=ParamType.COLUMN ) self.copy( up_proj_weight, hf_moe.shared_experts.up_proj.weight, param_type=ParamType.COLUMN ) self.copy( moe.shared_experts.linear_fc2.weight, hf_moe.shared_experts.down_proj.weight, param_type=ParamType.ROW ) 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(layer.input_layernorm.weight, hf_layer.input_layernorm.weight) else: self.set_selfattn_state(layer.self_attention, hf_layer.self_attn) self.copy(layer.self_attention.linear_qkv.layer_norm_weight, hf_layer.input_layernorm.weight) if hasattr(layer.mlp, 'router'): self.set_moe_layer_state(layer.mlp, hf_layer.mlp) self.copy(layer.pre_mlp_layernorm.weight, hf_layer.post_attention_layernorm.weight) else: self.set_mlp_state(layer.mlp, hf_layer.mlp) self.copy(layer.mlp.linear_fc1.layer_norm_weight, hf_layer.post_attention_layernorm.weight) def check_and_save(self, output_dir): sharded_info = split_torch_state_dict_into_shards( self._hfmodel.state_dict(), max_shard_size=self.max_shard_size ) global_shape = {self._hf_params_key_to_id[k]: v.shape for k, v in self._hfmodel.state_dict().items()} # select local bucket(s) for each rank n_savers = self.num_savers local_buckets = [] max_n_local_buckets = 1 if not sharded_info.is_sharded and self.rank == 0: local_buckets = list(sharded_info.filename_to_tensors.keys()) else: n_buckets = len(sharded_info.filename_to_tensors) rank = self.rank n_local_buckets = n_buckets // n_savers remainder = n_buckets % n_savers max_n_local_buckets = n_local_buckets + 1 if remainder == 0: start = rank * n_local_buckets max_n_local_buckets -= 1 elif rank < remainder: n_local_buckets += 1 start = rank * n_local_buckets else: start = rank * n_local_buckets + remainder local_buckets = list(sharded_info.filename_to_tensors.keys())[start:start + n_local_buckets] if rank == 0: index = { "metadata": sharded_info.metadata, "weight_map": sharded_info.tensor_to_filename, } os.makedirs(output_dir, exist_ok=True) with open(os.path.join(output_dir, "model.safetensors.index.json"), "w") as f: f.write(json.dumps(index, indent=2)) self._collect_dist_info() # In each iteration, all ranks save at most one local bucket for bucket_idx in range(max_n_local_buckets): required_keys = [] if bucket_idx < len(local_buckets): bucket_name = local_buckets[bucket_idx] required_keys: List[str] = sharded_info.filename_to_tensors[bucket_name] # build send/recv op across all ranks data, buffers, send_param_ids, recv_param_ids, ops = self._build_p2p_ops(required_keys) # run data sync if self.debug: logging.info(f"[Iters {bucket_idx} RANK {self.rank}] starts synchronizing parameters with other ranks...") if len(ops) > 0: reqs = dist.batch_isend_irecv(ops) if self.debug: for op in ops: if op.op == dist.isend: logging.info(f"[Iters {bucket_idx} RANK {self.rank}] ({self.rank} -> {op.peer}) with {op.tensor.numel() * op.tensor.dtype.itemsize / 2 ** 20} MiB.") else: logging.info(f"[Iters {bucket_idx} RANK {self.rank}] ({op.peer} -> {self.rank}) with {op.tensor.numel() * op.tensor.dtype.itemsize / 2 ** 20} MiB.") for req in reqs: req.wait() if self.debug: logging.info(f"[Iters {bucket_idx} RANK {self.rank}] finishes synchronizing") for remote_rank, param_ids in recv_param_ids.items(): for param_id, tensor in zip(param_ids, self._unpack_from_buffer(buffers[remote_rank], param_ids)): data[param_id][remote_rank] = tensor # apply merge function on the results # data: Dict[param_id, Dict[rank_id, tensor]] output_data = dict() for param_id, data_dict in data.items(): key = self._id_to_hf_params_key[param_id] param_type = ParamType(int(self._merge_type[param_id])) if param_type == ParamType.NULL: raise ValueError(f"ParamType.NULL found on {key}.") try: output_data[key] = self._merge_data(param_type, data_dict) except ParamMergeError as e: raise ValueError(f"Merge Error on key {key}: {e}") if output_data[key].shape != global_shape[param_id]: raise ValueError(f"Unexpected shape on {key}. Expected: {global_shape[param_id]}, but {output_data[key].shape}") del data # save safetensor files if bucket_idx < len(local_buckets): if not self.dryrun: safe_save_file( output_data, os.path.join(output_dir, local_buckets[bucket_idx]), metadata={"format": "pt"}, ) print(f"[Iters {bucket_idx} RANK {self.rank}] {local_buckets[bucket_idx]} is saved.") if self.debug: logging.debug(f"[Iters {bucket_idx} RANK {self.rank}] joined") dist.barrier() def _collect_dist_info(self): # Collect following metadatas: # param_id --> source_rank self._has_param = torch.zeros( [self.world_size, len(self._hf_params_key_to_id)], dtype=torch.bool, device=self.device ) for param_id in self._local_params.keys(): self._has_param[self.rank][param_id] = True dist.all_gather_into_tensor(self._has_param, self._has_param[self.rank]) self._has_param = self._has_param.T # param_id --> tensor_shape Dict[int, Tuple[int, ...]] # param_id --> tensor_dtype Dict[int, dtype] for param_id, param in self._local_params.items(): self._tensor_shape[param_id] = param.shape self._tensor_dtype[param_id] = param.dtype # collect across ranks tensor_shapes = [None] * self.world_size tensor_dtypes = [None] * self.world_size dist.all_gather_object(tensor_shapes, self._tensor_shape) dist.all_gather_object(tensor_dtypes, self._tensor_dtype) # NOTE: merge them together for rank, remote_shape in enumerate(tensor_shapes): if rank == self.rank: continue for remote_key, shape in remote_shape.items(): if remote_key not in self._tensor_shape: self._tensor_shape[remote_key] = shape elif shape != self._tensor_shape[remote_key]: raise ValueError( f"Find mismatched shape on local rank {self.rank} and remote rank {rank}, local shape: {self._tensor_shape[remote_key]}; remote shape: {shape}" ) for rank, remote_dtype in enumerate(tensor_dtypes): if rank == self.rank: continue for remote_key, dtype in remote_dtype.items(): if remote_key not in self._tensor_dtype: self._tensor_dtype[remote_key] = dtype elif dtype != self._tensor_dtype[remote_key]: raise ValueError( f"Find mismatched dtype on local rank {self.rank} and remote rank {rank}, local shape: {self._tensor_dtype[remote_key]}; remote shape: {dtype}" ) # merge_type global_merge_type = torch.zeros([self.world_size, self.hf_size], dtype=self._merge_type.dtype, device=self.device) dist.all_gather_into_tensor(global_merge_type, self._merge_type) for remote_rank_id, remote_merge_type in enumerate(global_merge_type): if self.debug: and_mask = torch.logical_and(remote_merge_type > 0, self._merge_type > 0) if (self._merge_type[and_mask] != remote_merge_type[and_mask]).any(): param_id = -1 for param_id in range(self.hf_size): if ( self._merge_type[param_id] > 0 and remote_merge_type[param_id] > 0 and self._merge_type[param_id] != remote_merge_type[param_id] ): break key = self._id_to_hf_params_key[param_id] raise ValueError(f"Find mismatched merge_type between local rank {self.rank} and remote rank {remote_rank_id} on key {key}") self._merge_type[remote_merge_type > 0] = remote_merge_type[remote_merge_type > 0] def _unpack_from_buffer(self, buffer: torch.Tensor, param_ids: List[int]) -> List[torch.Tensor]: start = 0 datas = [] for param_id in param_ids: shape = self._tensor_shape[param_id] dtype = self._tensor_dtype[param_id] offset = shape.numel() * dtype.itemsize datas.append(buffer[start:start + offset].view(dtype).view(shape).clone()) start += offset if start != buffer.numel(): raise ValueError(f"Expect {start} bytes from remote, but got {buffer.numel()} bytes!") return datas def _pack_into_byte_buffer(self, tensors: List[torch.Tensor]) -> torch.Tensor: return torch.cat([t.flatten().view(torch.uint8) for t in tensors]) def _build_p2p_ops(self, required_keys: List[str]): required_ids = torch.zeros( [self.world_size, self.hf_size], dtype=torch.bool, device=self.device ) for k in required_keys: required_ids[self.rank][self._hf_params_key_to_id[k]] = True dist.all_gather_into_tensor(required_ids, required_ids[self.rank]) send_ops = [] if self.debug: # (param_id, src_rank, dst_rank) send_recv_pattern = torch.zeros([self.hf_size, self.world_size, self.world_size], dtype=torch.int, device=self.device) send_param_ids = defaultdict(list) datas = defaultdict(list) # NOTE: for rank i to rank j, send params by ascending order # NOTE: to avoid hangs observed in multi-nodes scenarios, we merge tensors sent to same remote rank into a large uint8 tensor for param_id, has_data in enumerate(self._has_param[:, self.rank]): if not has_data: continue # group by remote_id for remote_rank, should_send in enumerate(required_ids[:, param_id]): if not should_send or remote_rank == self.rank: continue # NOTE: for each receiver, send param in ascending order by id data = self._local_params[param_id] if data.device != self.device: logging.warning(f"Find unexpected device {data.device} on key {self._id_to_hf_params_key[param_id]}, moving to {self.device}") data = data.to(self.device) if data.dtype != self._tensor_dtype[param_id]: raise ValueError(f"Get mismatched data type on key {self._id_to_hf_params_key[param_id]}") datas[remote_rank].append(data) send_param_ids[remote_rank].append(param_id) if self.debug: send_recv_pattern[param_id, self.rank, remote_rank] += 1 for remote_rank, raw_data in datas.items(): if len(raw_data) > 0: send_ops.append(dist.P2POp( dist.isend, self._pack_into_byte_buffer(raw_data), peer=remote_rank, tag=self.rank * self.world_size + remote_rank # (sender_rank, receiver_rank) ignored in NCCL )) recv_ops = [] collected_data = defaultdict(dict) buffer_size = [0] * self.world_size recv_param_ids = defaultdict(list) # NOTE: for rank i to rank j, recv params by ascending order for param_id, is_required in enumerate(required_ids[self.rank]): if not is_required: continue for remote_rank, has_data in enumerate(self._has_param[param_id]): if not has_data: continue if remote_rank == self.rank: collected_data[param_id][remote_rank] = self._local_params[param_id] else: recv_param_ids[remote_rank].append(param_id) shape = self._tensor_shape[param_id] dtype = self._tensor_dtype[param_id] buffer_size[remote_rank] += shape.numel() * dtype.itemsize if self.debug: send_recv_pattern[param_id, remote_rank, self.rank] -= 1 buffers = [None] * self.world_size for remote_rank, rank_size in enumerate(buffer_size): if rank_size == 0: continue buffers[remote_rank] = torch.empty(rank_size, dtype=torch.uint8, device=self.device) recv_ops.append(dist.P2POp( dist.irecv, buffers[remote_rank], peer=remote_rank, tag=remote_rank * self.world_size + self.rank # (sender_rank, receiver_rank) ignored in NCCL )) if self.debug: dist.all_reduce(send_recv_pattern) if send_recv_pattern.sum() != 0: for param_id, pattern_per_param in enumerate(send_recv_pattern): if pattern_per_param.sum() != 0: raise ValueError(f"Mismatched send/recv ops detected on key {self._id_to_hf_params_key[param_id]}: {pattern_per_param}.") raise ValueError("Mismatched send/recv ops detected.") logging.debug(f"[RANK {self.rank}] {len(send_ops)} send op & {len(recv_ops)} recv op.") return collected_data, buffers, send_param_ids, recv_param_ids, (send_ops + recv_ops) def _merge_data(self, merge_type: ParamType, tensor_dict) -> torch.Tensor: """Merge ShardedTensor collected across the group on this rank. If DP or EDP is larger than 1, tensor_dict could contains multiple data copies and an data deduplication isrequired. Args: merge_type (ParamType): Themerge policy be applied on the ShardedTensor tensor_dict (Dict[int, torch.Tensor]): The collected dict of ShardedTensors, mapping global_rank to tensor data. Returns: torch.Tensor: The merged tensor """ # tensor_dict: Dict[remote_rank, torch.Tensor] def merge_along_axis(axis, tensor_dict, is_expert: bool = False): global_ranks = torch.tensor(list(tensor_dict.keys()), dtype=torch.long, device=self.device) # (N, ) ranks = self._rank_mapping.index_select(0, global_ranks) # (N, 6) if self.debug: if is_expert and ranks[:, 5].any(): raise ParamMergeError("Unexpected expert parameter data from non-zero expert dp rank") if not is_expert and ranks[:, 4].any(): raise ParamMergeError("Unexpected parameter data from non-zero dp rank") if ranks[:, 1].max() != ranks[:, 1].min(): raise ParamMergeError("Unexpected parameter data from multiple pp ranks") if is_expert and ranks[:, 3].max() != ranks[:, 3].min(): raise ParamMergeError("Unexpected expert parameter data from multiple ep ranks") def deduplicate_and_sort(tensor_dict, rank_group: int): tensors = dict() for global_rank, data in tensor_dict.items(): rank = self._rank_mapping[global_rank][rank_group] tensors[int(rank)] = data return [item[1] for item in sorted(tensors.items())] tensors = list(tensor_dict.values())[:1] if is_expert: tensors = deduplicate_and_sort(tensor_dict, 2) else: tensors = deduplicate_and_sort(tensor_dict, 0) return torch.cat(tensors, dim=axis) def no_merge_func(tensor_dict): return list(tensor_dict.values())[0] def merge_qkv(is_bias, tensor_dict): res = merge_along_axis(0, tensor_dict) if is_bias: assert res.shape[-1] == 1 return res.flatten() return res.flatten(0, 1) merge_func_mapping = { ParamType.MOE_COLUMN: partial(merge_along_axis, 0, is_expert=True), ParamType.MOE_ROW: partial(merge_along_axis, 1, is_expert=True), ParamType.COLUMN: partial(merge_along_axis, 0), ParamType.ROW: partial(merge_along_axis, 1), ParamType.QKV_W: partial(merge_qkv, False), ParamType.QKV_B: partial(merge_qkv, True), ParamType.UNIQUE: no_merge_func, } return merge_func_mapping[merge_type](tensor_dict)