# coding=utf-8 # Copyright 2023 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import json import os from functools import partial from pathlib import Path from typing import Any, Callable, Dict, List, Literal, Optional, Union import torch from huggingface_hub import split_torch_state_dict_into_shards from transformers.utils import ( SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME, ) from ...utils.import_utils import is_peft_available from ...utils.require_utils import requires_neuronx_distributed, requires_safetensors, requires_torch_xla from .modeling_utils import MODEL_PARALLEL_SHARDS_DIR_NAME from .transformations_utils import ModelWeightTransformationSpecs, to_original_weights if is_peft_available(): from peft.utils.constants import ( SAFETENSORS_WEIGHTS_NAME as PEFT_SAFETENSORS_WEIGHTS_NAME, ) from peft.utils.constants import ( WEIGHTS_NAME as PEFT_WEIGHTS_NAME, ) else: PEFT_SAFETENSORS_WEIGHTS_NAME = PEFT_WEIGHTS_NAME = "" @requires_torch_xla def xser_load_on_cpu(path: str): """ Modified version from neuronx_distributed `_xser_load` function load located at: https://github.com/aws-neuron/neuronx-distributed/blob/e83494557cb4c5b7e185ccf6c9240bfed9a1993d/src/neuronx_distributed/parallel_layers/checkpointing.py#L252 Instead of moving the loaded tensors to the XLA device, it keeps them on CPU. """ import torch_xla.core.xla_model as xm import torch_xla.utils.serialization as xser ref_data = torch.load(path) def convert_fn(tensors): rewritten_tensors = [] for t in tensors: rewritten_tensors.append(torch.load(os.path.join(path + ".tensors", "tensor_{}.pt".format(t.tid)))) return rewritten_tensors def select_fn(v): return type(v) is xser.TensorReference return xm.ToXlaTensorArena(convert_fn, select_fn).transform(ref_data) def consolidate_tensor_parallel_checkpoints( sharded_checkpoints: List[Path], load_function: Callable[[Union[str, Path]], Dict[str, Any]], metadata: Dict[str, Any], adapter_name: Optional[str] = None, ) -> Dict[str, "torch.Tensor"]: state_dicts = [] sharded_checkpoints = sorted(sharded_checkpoints) for sharded_checkpoint in sharded_checkpoints: if not sharded_checkpoint.is_file(): continue state_dicts.append(load_function(sharded_checkpoint.as_posix())) parameters_metadata = metadata["parameters"] transformation_specs_metadata = metadata["model_weight_transformation_specs"] # We recreate the transformation specs from the metadata. transformations_specs = [] for specs_metadata in transformation_specs_metadata: specs = ModelWeightTransformationSpecs.from_metadata(specs_metadata) transformations_specs.append(specs) # We transform the sharded state dicts as follows: # [state_dict_tp_rank_0, state_dict_tp_rank_1, ...] # -> { # key: [state_dict_tp_rank_0[key], state_dict_tp_rank_1[key], ...], # for key in state_dict_tp_rank_0.keys() # } parameter_names = state_dicts[0].keys() sharded_state_dicts = {name: [state_dict[name] for state_dict in state_dicts] for name in parameter_names} consolidated_state_dict = to_original_weights( transformations_specs, sharded_state_dicts, parameters_metadata, adapter_name=adapter_name ) return consolidated_state_dict @requires_neuronx_distributed def consolidate_model_parallel_checkpoints( checkpoint_dir: Path, adapter_name: Optional[str] = None ) -> Dict[str, "torch.Tensor"]: model_checkpoint_dir = checkpoint_dir / "model" # Case 1: the checkpoint was saved with xser. sharded_checkpoints = list(model_checkpoint_dir.glob("dp_rank*.tensors")) if sharded_checkpoints: sharded_checkpoints = model_checkpoint_dir.glob("dp_rank_*") sharded_checkpoints = [ p for p in sharded_checkpoints if not (p.name.endswith("info.pt") or p.name.endswith("tensors")) ] load_function = xser_load_on_cpu # Case 2: If no file was found, maybe the checkpoint was saved without xser. if not sharded_checkpoints: sharded_checkpoints = list(model_checkpoint_dir.glob("dp_rank_*.pt")) load_function = partial(torch.load, weights_only=True) if not sharded_checkpoints: raise ValueError(f"Could not find any sharded checkpoint in {model_checkpoint_dir.as_posix()}") pp_size = max((int(checkpoint_path.stem[-2:]) for checkpoint_path in sharded_checkpoints)) + 1 checkpoints_grouped_by_pp_ranks = [[] for _ in range(pp_size)] metadatas = [] for pp_rank in range(pp_size): for checkpoint_path in sharded_checkpoints: checkpoint_name = checkpoint_path.stem if int(checkpoint_name[-2:]) == pp_rank: checkpoints_grouped_by_pp_ranks[pp_rank].append(checkpoint_path) if (checkpoint_dir / f"mp_metadata_pp_rank_{pp_rank}.pt").is_file(): metadatas.append(torch.load(checkpoint_dir / f"mp_metadata_pp_rank_{pp_rank}.pt")) else: with open(checkpoint_dir / f"mp_metadata_pp_rank_{pp_rank}.json") as fp: metadatas.append(json.load(fp)) consolidated_state_dict = {} for pp_rank, checkpoint_group_for_pp_rank in enumerate(checkpoints_grouped_by_pp_ranks): consolidated_for_pp_rank = consolidate_tensor_parallel_checkpoints( checkpoint_group_for_pp_rank, load_function, metadatas[pp_rank], adapter_name=adapter_name, ) consolidated_state_dict.update(**consolidated_for_pp_rank) for key, tensor in consolidated_state_dict.items(): consolidated_state_dict[key] = tensor return consolidated_state_dict @requires_safetensors def consolidate_model_parallel_checkpoints_to_unified_checkpoint( checkpoint_dir: Union[str, Path], output_dir: Union[str, Path], save_format: Literal["pytorch", "safetensors"] = "safetensors", ): from safetensors.torch import save_file # We import here to avoid circular import. from ...peft.peft_model import ADAPTER_MODEL_PARALLEL_SHARDS_DIR_NAME if not isinstance(checkpoint_dir, Path): checkpoint_dir = Path(checkpoint_dir) directories = list(checkpoint_dir.iterdir()) directories_to_consolidate = [] if checkpoint_dir.name != MODEL_PARALLEL_SHARDS_DIR_NAME: if (checkpoint_dir / MODEL_PARALLEL_SHARDS_DIR_NAME).is_dir(): directories_to_consolidate = [checkpoint_dir / MODEL_PARALLEL_SHARDS_DIR_NAME] else: for dir in directories: if dir.is_dir() and dir.name.startswith("adapter_"): directories_to_consolidate.append(dir / ADAPTER_MODEL_PARALLEL_SHARDS_DIR_NAME) if not directories_to_consolidate: raise ValueError(f"Could not find the tensor parallel shards from {checkpoint_dir}") else: directories_to_consolidate = [checkpoint_dir] if not isinstance(output_dir, Path): output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) for checkpoint_dir in directories_to_consolidate: # We need to go one level up because the checkpoint directory is at the shards level here. parent_dir = checkpoint_dir.parent current_output_dir = output_dir is_adapter_model = parent_dir.name.startswith("adapter_") adapter_name = None if is_adapter_model: safe_weights_name = PEFT_SAFETENSORS_WEIGHTS_NAME weights_name = PEFT_WEIGHTS_NAME if parent_dir.name != "adapter_default": adapter_name = parent_dir.name.split("_", maxsplit=1)[-1] current_output_dir = output_dir / adapter_name else: adapter_name = "default" else: safe_weights_name = SAFE_WEIGHTS_NAME weights_name = WEIGHTS_NAME current_output_dir.mkdir(parents=True, exist_ok=True) state_dict = consolidate_model_parallel_checkpoints(checkpoint_dir, adapter_name=adapter_name) state_dict_split = split_torch_state_dict_into_shards( state_dict, filename_pattern=safe_weights_name if save_format == "safetensors" else weights_name ) # Save index if sharded if state_dict_split.is_sharded: index = { "metadata": state_dict_split.metadata, "weight_map": state_dict_split.tensor_to_filename, } save_index_file = SAFE_WEIGHTS_INDEX_NAME if save_format == "safetensors" else WEIGHTS_INDEX_NAME with open(current_output_dir / save_index_file, "w") as fp: content = json.dumps(index, indent=2, sort_keys=True) + "\n" fp.write(content) # Save the model filename_to_tensors = state_dict_split.filename_to_tensors.items() for shard_file, tensors in filename_to_tensors: shard = {} for tensor in tensors: shard[tensor] = state_dict[tensor].contiguous() del state_dict[tensor] if save_format == "safetensors": save_file(shard, current_output_dir / shard_file, metadata={"format": "pt"}) else: torch.save(shard, current_output_dir / shard_file)