checkpoint_converter_fsdp_hf.py 2.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465
  1. # Copyright (c) Meta Platforms, Inc. and affiliates.
  2. # This software may be used and distributed according to the terms of the Llama 2 Community License Agreement.
  3. # from accelerate import init_empty_weights, load_checkpoint_and_dispatch
  4. import fire
  5. import os
  6. import sys
  7. import yaml
  8. from transformers import LlamaTokenizer
  9. from .model_utils import load_llama_from_config
  10. # Get the current file's directory
  11. current_directory = os.path.dirname(os.path.abspath(__file__))
  12. # Get the parent directory
  13. parent_directory = os.path.dirname(current_directory)
  14. # Append the parent directory to sys.path
  15. sys.path.append(parent_directory)
  16. from model_checkpointing import load_sharded_model_single_gpu
  17. def main(
  18. fsdp_checkpoint_path="", # Path to FSDP Sharded model checkpoints
  19. consolidated_model_path="", # Path to save the HF converted model checkpoints
  20. HF_model_path_or_name="" # Path/ name of the HF model that include config.json and tokenizer_config.json (e.g. meta-llama/Llama-2-7b-chat-hf)
  21. ):
  22. try:
  23. file_name = 'train_params.yaml'
  24. # Combine the directory and file name to create the full path
  25. train_params_path = os.path.join(fsdp_checkpoint_path, file_name)
  26. # Open the file
  27. with open(train_params_path, 'r') as file:
  28. # Load the YAML data
  29. data = yaml.safe_load(file)
  30. # Access the 'model_name' field
  31. HF_model_path_or_name = data.get('model_name')
  32. print(f"Model name: {HF_model_path_or_name}")
  33. except FileNotFoundError:
  34. print(f"The file {train_params_path} does not exist.")
  35. HF_model_path_or_name = input("Please enter the model name: ")
  36. print(f"Model name: {HF_model_path_or_name}")
  37. except Exception as e:
  38. print(f"An error occurred: {e}")
  39. #load the HF model definition from config
  40. model_def = load_llama_from_config(HF_model_path_or_name)
  41. print("model is loaded from config")
  42. #load the FSDP sharded checkpoints into the model
  43. model = load_sharded_model_single_gpu(model_def, fsdp_checkpoint_path)
  44. print("model is loaded from FSDP checkpoints")
  45. #loading the tokenizer form the model_path
  46. tokenizer = LlamaTokenizer.from_pretrained(HF_model_path_or_name)
  47. tokenizer.save_pretrained(consolidated_model_path)
  48. #save the FSDP sharded checkpoints in HF format
  49. model.save_pretrained(consolidated_model_path)
  50. print(f"HuggingFace model checkpoints has been saved in {consolidated_model_path}")
  51. if __name__ == "__main__":
  52. fire.Fire(main)