Skip to content

feat: add save_model_dir flag where final checkpoint saved #291

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Ssukriti merged 12 commits into from
Aug 14, 2024
Merged

feat: add save_model_dir flag where final checkpoint saved #291

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
b300216
add save_model_dir flag for final checkpoint
anhuong Aug 6, 2024
cee43b9
remove output_dir logic, add save method
anhuong Aug 7, 2024
78f25aa
update accelerate_launch, remove save tokenizer
anhuong Aug 8, 2024
b13333f
fix: put back creation of .complete file
anhuong Aug 8, 2024
0c8b6e5
fix failing tests and add new ones
anhuong Aug 8, 2024
b8fc646
Merge branch 'main' into save_model_dir
anhuong Aug 8, 2024
2a9197c
tests: add sft_trainer test to train and save
anhuong Aug 8, 2024
d28ec4d
add docs on saving checkpoints and fix help msg
anhuong Aug 8, 2024
536442c
update example and note best checkpoint
anhuong Aug 9, 2024
9efc549
merge changes from main branch
anhuong Aug 13, 2024
59f840f
changes based on PR review
anhuong Aug 13, 2024
9954678
add logging to save, fix error out properly
anhuong Aug 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
227 changes: 95 additions & 132 deletions build/accelerate_launch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,6 @@
import subprocess
import sys
import traceback
import tempfile
import shutil
from pathlib import Path
import json

Expand All @@ -37,12 +35,9 @@
# Local
from build.utils import (
process_accelerate_launch_args,
serialize_args,
get_highest_checkpoint,
copy_checkpoint,
)
from tuning.utils.config_utils import get_json_config
from tuning.config.tracker_configs import FileLoggingTrackerConfig
from tuning.utils.error_logging import (
write_termination_log,
USER_ERROR_EXIT_CODE,
Expand Down Expand Up @@ -102,143 +97,111 @@ def main():
# Launch training
#
##########
original_output_dir = job_config.get("output_dir")
with tempfile.TemporaryDirectory() as tempdir:
try:
# checkpoints outputted to tempdir, only final checkpoint copied to output dir
job_config["output_dir"] = tempdir
updated_args = serialize_args(job_config)
os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = updated_args

launch_command(args)
except subprocess.CalledProcessError as e:
# If the subprocess throws an exception, the base exception is hidden in the
# subprocess call and is difficult to access at this level. However, that is not
# an issue because sft_trainer.py would have already written the exception
# message to termination log.
logging.error(traceback.format_exc())
# The exit code that sft_trainer.py threw is captured in e.returncode

return_code = e.returncode
if return_code not in [INTERNAL_ERROR_EXIT_CODE, USER_ERROR_EXIT_CODE]:
return_code = INTERNAL_ERROR_EXIT_CODE
write_termination_log(f"Unhandled exception during training. {e}")
sys.exit(return_code)
except Exception as e: # pylint: disable=broad-except
logging.error(traceback.format_exc())
output_dir = job_config.get("output_dir")
try:
# checkpoints outputted to tempdir, only final checkpoint copied to output dir
launch_command(args)
except subprocess.CalledProcessError as e:
# If the subprocess throws an exception, the base exception is hidden in the
# subprocess call and is difficult to access at this level. However, that is not
# an issue because sft_trainer.py would have already written the exception
# message to termination log.
logging.error(traceback.format_exc())
# The exit code that sft_trainer.py threw is captured in e.returncode

return_code = e.returncode
if return_code not in [INTERNAL_ERROR_EXIT_CODE, USER_ERROR_EXIT_CODE]:
return_code = INTERNAL_ERROR_EXIT_CODE
write_termination_log(f"Unhandled exception during training. {e}")
sys.exit(INTERNAL_ERROR_EXIT_CODE)
sys.exit(return_code)
except Exception as e: # pylint: disable=broad-except
logging.error(traceback.format_exc())
write_termination_log(f"Unhandled exception during training. {e}")
sys.exit(INTERNAL_ERROR_EXIT_CODE)

try:
last_checkpoint_dir = get_highest_checkpoint(tempdir)
last_checkpoint_path = os.path.join(tempdir, last_checkpoint_dir)
# remove lm_head from granite with llama arch models
try:
checkpoint_dir = job_config.get("save_model_dir")
if not checkpoint_dir:
checkpoint_dir = os.path.join(
output_dir, get_highest_checkpoint(output_dir)

@ashokponkumar ashokponkumar Aug 9, 2024

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we know if removing the lm_head is required while resuming a training, if so we should think about removing the lm_head from each intermediate checkpoint as it is being written.

@anhuong anhuong Aug 9, 2024

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removing lm_head is required for loading the tuned model with vLLM, but I don't think it's required for resuming training. Is there a way I could check resuming training? If you have a command to run with sft_trainer, that would be super helpful!

@ashokponkumar ashokponkumar Aug 9, 2024

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we have to stop a run midway and then run the training again.

We have to ensure we use https://huggingface.co/docs/transformers/en/main_classes/trainer#transformers.Trainer.train.resume_from_checkpoint argument. Seems like we are not currently setting it. We should possibly set it. This will resume from the previously stopped last checkpoint.

@Ssukriti Ssukriti Aug 10, 2024

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hi @ashokponkumar we have not heard of any requirement to remove lm_head to resume training. @Abhishek-TAMU will be working on resume training from checkpoint in follow up PR as per issue https://github.ibm.com/ai-foundation/watson-fm-stack-tracker/issues/1007 . Can we wait for him to add the change?

ypu can post any tips on the issue @ashokponkumar

@ashokponkumar ashokponkumar Aug 10, 2024

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure Sukriti.

)

use_flash_attn = job_config.get("use_flash_attn", True)
adapter_config_path = os.path.join(checkpoint_dir, "adapter_config.json")
tokenizer = AutoTokenizer.from_pretrained(checkpoint_dir)

use_flash_attn = job_config.get("use_flash_attn", True)
adapter_config_path = os.path.join(
last_checkpoint_path, "adapter_config.json"
if os.path.exists(adapter_config_path):
base_model_path = get_base_model_from_adapter_config(adapter_config_path)
base_model = AutoModelForCausalLM.from_pretrained(
base_model_path,
attn_implementation="flash_attention_2" if use_flash_attn else None,
torch_dtype=bfloat16 if use_flash_attn else None,
)
tokenizer = AutoTokenizer.from_pretrained(last_checkpoint_path)

if os.path.exists(adapter_config_path):
base_model_path = get_base_model_from_adapter_config(
adapter_config_path
)
base_model = AutoModelForCausalLM.from_pretrained(
base_model_path,
attn_implementation="flash_attention_2" if use_flash_attn else None,
torch_dtype=bfloat16 if use_flash_attn else None,
)
# since the peft library (PEFTModelForCausalLM) does not handle cases
# where the model's layers are modified, in our case the embedding layer
# is modified, so we resize the backbone model's embedding layer with our own
# utility before passing it along to load the PEFT model.
tokenizer_data_utils.tokenizer_and_embedding_resize(
{}, tokenizer=tokenizer, model=base_model
)
model = PeftModel.from_pretrained(
base_model,
checkpoint_dir,
attn_implementation="flash_attention_2" if use_flash_attn else None,
torch_dtype=bfloat16 if use_flash_attn else None,
)
else:
model = AutoModelForCausalLM.from_pretrained(
checkpoint_dir,
attn_implementation="flash_attention_2" if use_flash_attn else None,
torch_dtype=bfloat16 if use_flash_attn else None,
)

# since the peft library (PEFTModelForCausalLM) does not handle cases
# where the model's layers are modified, in our case the embedding layer
# is modified, so we resize the backbone model's embedding layer with our own
# utility before passing it along to load the PEFT model.
tokenizer_data_utils.tokenizer_and_embedding_resize(
{}, tokenizer=tokenizer, model=base_model
model_arch = model.config.model_type
# check that it is a granite model with llama architecture with tied weights
# ie. lm_head is duplicate of embeddings

# a fine tuned model will have params_dict.get("model.embed_tokens.weight")
# a prompt adapter has params_dict.get("base_model.model.embed_tokens.weight")
# a lora adapter has params_dict.get("base_model.model.model.embed_tokens.weight")
if model_arch == "llama" and hasattr(model, "lm_head"):
if (
# lora tuned model has an addt model layer
(
hasattr(model.model, "model")
and model.lm_head.weight.untyped_storage().data_ptr()
== model.model.model.embed_tokens.weight.untyped_storage().data_ptr()
)
model = PeftModel.from_pretrained(
base_model,
last_checkpoint_path,
attn_implementation="flash_attention_2" if use_flash_attn else None,
torch_dtype=bfloat16 if use_flash_attn else None,
)
else:
model = AutoModelForCausalLM.from_pretrained(
last_checkpoint_path,
attn_implementation="flash_attention_2" if use_flash_attn else None,
torch_dtype=bfloat16 if use_flash_attn else None,
# prompt tuned model or fine tuned model
or (
hasattr(model.model, "embed_tokens")
and model.lm_head.weight.untyped_storage().data_ptr()
== model.model.embed_tokens.weight.untyped_storage().data_ptr()
)
):

model_arch = model.config.model_type
# check that it is a granite model with llama architecture with tied weights
# ie. lm_head is duplicate of embeddings

# a fine tuned model will have params_dict.get("model.embed_tokens.weight")
# a prompt adapter has params_dict.get("base_model.model.embed_tokens.weight")
# a lora adapter has params_dict.get("base_model.model.model.embed_tokens.weight")
copy_checkpoint_bool = True
if model_arch == "llama" and hasattr(model, "lm_head"):
if (
# lora tuned model has an addt model layer
(
hasattr(model.model, "model")
and model.lm_head.weight.untyped_storage().data_ptr()
== model.model.model.embed_tokens.weight.untyped_storage().data_ptr()
)
# prompt tuned model or fine tuned model
or (
hasattr(model.model, "embed_tokens")
and model.lm_head.weight.untyped_storage().data_ptr()
== model.model.embed_tokens.weight.untyped_storage().data_ptr()
)
):

copy_checkpoint_bool = False
logging.info("Removing lm_head from checkpoint")
del model.lm_head.weight

if hasattr(model, "lm_head.weight"):
logging.warning("Failed to delete lm_head.weight from model")

logging.info("Saving checkpoint to %s", original_output_dir)
model.save_pretrained(original_output_dir)
# save tokenizer with model
tokenizer.save_pretrained(original_output_dir)

# copy last checkpoint into mounted output dir
if copy_checkpoint_bool:
logging.info(
"Copying last checkpoint %s into output dir %s",
last_checkpoint_dir,
original_output_dir,
)
copy_checkpoint(last_checkpoint_path, original_output_dir)
except Exception as e: # pylint: disable=broad-except
logging.error(traceback.format_exc())
write_termination_log(
f"Exception encountered writing output model to storage: {e}"
)
sys.exit(INTERNAL_ERROR_EXIT_CODE)
logging.info("Removing lm_head from checkpoint")
del model.lm_head.weight

# copy over any loss logs
try:
train_logs_filepath = os.path.join(
tempdir,
FileLoggingTrackerConfig.training_logs_filename,
)
if os.path.exists(train_logs_filepath):
shutil.copy(train_logs_filepath, original_output_dir)

# The .complete file will signal to users that we are finished copying
# files over
if os.path.exists(original_output_dir):
Path(os.path.join(original_output_dir, ".complete")).touch()
except Exception as e: # pylint: disable=broad-except
logging.error(traceback.format_exc())
write_termination_log(
f"Exception encountered in capturing training logs: {e}"
)
sys.exit(INTERNAL_ERROR_EXIT_CODE)
if hasattr(model, "lm_head.weight"):
logging.warning("Failed to delete lm_head.weight from model")

logging.info("Saving checkpoint to %s", output_dir)
model.save_pretrained(checkpoint_dir)
# save tokenizer with model
tokenizer.save_pretrained(checkpoint_dir)

except Exception as e: # pylint: disable=broad-except
logging.error(traceback.format_exc())
write_termination_log(f"Exception encountered removing lm_head from model: {e}")
sys.exit(INTERNAL_ERROR_EXIT_CODE)

# The .complete file will signal to users that we are finished copying
# files over
if os.path.exists(output_dir):
Path(os.path.join(output_dir, ".complete")).touch()

return 0

Expand Down
Loading