From 02b34c2bf1a2011a24f0e55c31fd0d79410cb041 Mon Sep 17 00:00:00 2001 From: EeyoreLee Date: Fri, 4 Aug 2023 15:45:40 +0800 Subject: [PATCH] zero_to_fp32 script adds support for tag argument --- deepspeed/utils/zero_to_fp32.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/deepspeed/utils/zero_to_fp32.py b/deepspeed/utils/zero_to_fp32.py index fd5a2eb8bfdd..c98caae31534 100755 --- a/deepspeed/utils/zero_to_fp32.py +++ b/deepspeed/utils/zero_to_fp32.py @@ -574,9 +574,14 @@ def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None): "output_file", type=str, help="path to the pytorch fp32 state_dict output file (e.g. path/checkpoint-12/pytorch_model.bin)") + parser.add_argument("-t", + "--tag", + type=str, + default=None, + help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1") parser.add_argument("-d", "--debug", action='store_true', help="enable debug") args = parser.parse_args() debug = args.debug - convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, args.output_file) + convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, args.output_file, tag=args.tag)