From b0ca64d2959dba652bd911fa5f810ebee441002d Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 25 Jul 2024 20:07:37 -0700 Subject: [PATCH 1/3] add --- llmfoundry/callbacks/hf_checkpointer.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py index 35508cc0c7..946227ff34 100644 --- a/llmfoundry/callbacks/hf_checkpointer.py +++ b/llmfoundry/callbacks/hf_checkpointer.py @@ -375,6 +375,17 @@ def transform_config( if 'moe_world_size' in getattr(copied_config, 'ffn_config', {}): copied_config.ffn_config['moe_world_size'] = 1 return copied_config + + def pre_register_edit(self, local_save_path: str): + """Edit the model before registering with MLflow. + + This allows a subclass to modify the model before registering with MLflow. The base class implementation will + make no modifications. + + Args: + local_save_path (str): The path to the model to be transformed. + """ + pass def transform_model_pre_registration( self, @@ -618,6 +629,10 @@ def tensor_hook( os.path.join(local_save_path, license_filename), ) + self.pre_register_edit( + local_save_path, + ) + # Spawn a new process to register the model. process = SpawnProcess( target=_register_model_with_run_id_multiprocess, From 020842d2cd1143d245302d1a21fc62e6f38954f9 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 26 Jul 2024 12:13:56 -0700 Subject: [PATCH 2/3] pc --- llmfoundry/callbacks/hf_checkpointer.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py index 946227ff34..1797c3b5b4 100644 --- a/llmfoundry/callbacks/hf_checkpointer.py +++ b/llmfoundry/callbacks/hf_checkpointer.py @@ -375,7 +375,7 @@ def transform_config( if 'moe_world_size' in getattr(copied_config, 'ffn_config', {}): copied_config.ffn_config['moe_world_size'] = 1 return copied_config - + def pre_register_edit(self, local_save_path: str): """Edit the model before registering with MLflow. @@ -629,9 +629,7 @@ def tensor_hook( os.path.join(local_save_path, license_filename), ) - self.pre_register_edit( - local_save_path, - ) + self.pre_register_edit(local_save_path,) # Spawn a new process to register the model. process = SpawnProcess( From d1a54657c2d35cf25856f94cc195d05055a45fec Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 26 Jul 2024 12:16:04 -0700 Subject: [PATCH 3/3] add test --- tests/a_scripts/inference/test_convert_composer_to_hf.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/a_scripts/inference/test_convert_composer_to_hf.py b/tests/a_scripts/inference/test_convert_composer_to_hf.py index ffdb09ca98..9eb214e83d 100644 --- a/tests/a_scripts/inference/test_convert_composer_to_hf.py +++ b/tests/a_scripts/inference/test_convert_composer_to_hf.py @@ -388,6 +388,9 @@ def test_huggingface_conversion_callback_interval( checkpointer_callback.transform_model_pre_registration = MagicMock( wraps=checkpointer_callback.transform_model_pre_registration, ) + checkpointer_callback.pre_register_edit = MagicMock( + wraps=checkpointer_callback.pre_register_edit, + ) trainer = Trainer( model=original_model, device='gpu', @@ -413,9 +416,11 @@ def test_huggingface_conversion_callback_interval( metadata={}, ) assert checkpointer_callback.transform_model_pre_registration.call_count == 1 + assert checkpointer_callback.pre_register_edit.call_count == 1 assert mlflow_logger_mock.register_model_with_run_id.call_count == 1 else: assert checkpointer_callback.transform_model_pre_registration.call_count == 0 + assert checkpointer_callback.pre_register_edit.call_count == 0 assert mlflow_logger_mock.save_model.call_count == 0 assert mlflow_logger_mock.register_model_with_run_id.call_count == 0