| import os |
| import torch |
| import time |
|
|
| from importlib import import_module |
| from functools import partial |
|
|
| |
| |
| from megatron.training import global_vars |
| global_vars.build_tokenizer = partial(global_vars.build_tokenizer, trust_remote_code=True) |
| from megatron.training.initialize import initialize_megatron |
| from megatron.training import get_args |
|
|
| @torch.inference_mode() |
| def convert( |
| synchronizer, |
| pretrain_script, |
| load_dir, |
| save_dir, |
| hf_dir: str = None, |
| mcore2hf: bool = False |
| ): |
| sync_module = import_module(synchronizer) |
| model_provider_func = None |
| if pretrain_script is not None: |
| pretrain_module = import_module(pretrain_script) |
| model_provider_func = pretrain_module.model_provider |
| if mcore2hf: |
| args = get_args() |
| args.load = load_dir |
| synchronizer = sync_module.MG2HFSynchronizer( |
| hf_dir, |
| model_provider_func=model_provider_func |
| ) |
| else: |
| synchronizer = sync_module.HF2MGSynchronizer( |
| load_dir, |
| model_provider_func=model_provider_func |
| ) |
|
|
| synchronizer.sync_params() |
| synchronizer.check_and_save(save_dir) |
|
|
| def add_args(parser): |
| group = parser.add_argument_group(title='Distributed CKPT Convertor') |
|
|
| group.add_argument('--synchronizer', type=str, default='general', |
| help='The module required for conversion, should have two class `MG2HFSynchronizer` and `HF2MGSynchronizer`') |
| group.add_argument('--pretrain-script', type=str, default=None, |
| help='The path of pretrain script, the model is built with the `model_provider()` in it, use `pretrain_gpt.py` by default.') |
| group.add_argument('--model-type', type=str, required=True, choices=['GPT']) |
| group.add_argument('--load-dir', type=str, required=True, help='Directory to load model checkpoint from') |
| group.add_argument('--save-dir', type=str, required=True, help='Directory to save model checkpoint to') |
| group.add_argument('--hf-dir', type=str, help='pretrained huggingface checkpoint directory') |
|
|
| group.add_argument('--padded-vocab-size', type=int, default=None) |
| group.add_argument('--target-ckpt-format', default='torch_dist', choices=['torch_dist'], help='Checkpoint format to use.') |
| group.add_argument('--use-gpu', action='store_true') |
| group.add_argument('--mcore2hf', action='store_true') |
| group.add_argument('--debug', action='store_true', help='enable debug mode to check the integrity of conversion, could be slow for large models') |
| group.add_argument('--dryrun', action='store_true', help='If set, the conversion will be performed on an empty model and no save occurs, only for mcore2hf debugging.') |
| group.add_argument('--num-hf-saver', type=int, default=None, help='Set the amount of huggingface savers in mcore2hf mode, each saver will requires extra memorys for comm, set smaller if OOM occurs, by default it is world_size.') |
| group.add_argument('--max-shard-size', type=str, default='4GB', help='Set the sharded size, reduce memory consumption if smaller') |
| |
| return parser |
|
|
| if __name__ == '__main__': |
| start_time = time.time() |
| initialize_megatron(extra_args_provider=add_args, allow_no_cuda=True) |
| args = get_args() |
| actual_world_size = int(os.environ["WORLD_SIZE"]) |
| actual_rank = int(os.environ["RANK"]) |
| if args.num_hf_saver is None: |
| args.num_hf_saver = actual_world_size |
| print(f"WORLD_SIZE: {actual_world_size}, RANK: {actual_rank}, LOCAL_RANK: {args.local_rank}") |
|
|
| if args.target_ckpt_format == 'torch_dist': |
| convert( |
| args.synchronizer, |
| args.pretrain_script, |
| args.load_dir, |
| args.save_dir, |
| hf_dir=args.hf_dir, |
| mcore2hf=args.mcore2hf |
| ) |
| else: |
| raise NotImplementedError() |
|
|
| torch.distributed.barrier() |
| end_time = time.time() |
| if actual_rank == 0: |
| print(f"Conversion finished in {end_time - start_time} seconds.") |
