From 123c3563130eecaa213f135aa620d655e5e9063b Mon Sep 17 00:00:00 2001 From: zhangir-azerbayev Date: Tue, 24 Oct 2023 15:32:22 -0600 Subject: [PATCH] change behavior of finetuning argument --- megatron/checkpointing.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/megatron/checkpointing.py b/megatron/checkpointing.py index 8bcc01f3b..ea447faa5 100644 --- a/megatron/checkpointing.py +++ b/megatron/checkpointing.py @@ -351,20 +351,18 @@ def load_checkpoint( ): """Load a model checkpoint and return the iteration.""" if neox_args.deepspeed: - load_optim_and_scheduler = ( + load_optim= ( not neox_args.no_load_optim ) # TODO: These should be configured by separate args - if neox_args.finetune: - load_optim_and_scheduler = False if iteration is not None: tag = get_checkpoint_tag(iteration) else: tag = None checkpoint_name, state_dict = model.load_checkpoint( neox_args.load, - load_optimizer_states=load_optim_and_scheduler, - load_lr_scheduler_states=load_optim_and_scheduler, - load_module_only=not load_optim_and_scheduler, + load_optimizer_states=load_optim, + load_lr_scheduler_states=False, + load_module_only=not load_optim, tag=tag, )