merge main

src/llmcompressor/__init__.py

@@ -38,8 +38,6 @@
     active_session,
     callbacks,
     create_session,
-    finalize,
-    initialize,
     reset_session,
 )
-from llmcompressor.entrypoints import Oneshot, oneshot
+from llmcompressor.entrypoints import Oneshot, oneshot, train

src/llmcompressor/core/__init__.py

@@ -13,8 +13,6 @@
     active_session,
     callbacks,
     create_session,
-    finalize,
-    initialize,
     reset_session,
 )
 from llmcompressor.core.state import Data, Hardware, ModifiedState, State
@@ -35,8 +33,6 @@
     "create_session",
     "active_session",
     "reset_session",
-    "initialize",
-    "finalize",
     "apply",
     "callbacks",
     "LifecycleCallbacks",

src/llmcompressor/transformers/finetune/session_mixin.py

@@ -11,7 +11,7 @@
 from transformers.trainer_callback import TrainerState
 from transformers.trainer_utils import get_last_checkpoint
 
-from llmcompressor.core import active_session, callbacks, create_session, finalize
+from llmcompressor.core import active_session, callbacks, create_session
 from llmcompressor.metrics import LoggerManager
 from llmcompressor.modifiers.distillation.utils.pytorch.model_wrapper import (
     KDModelWrapper,
@@ -182,7 +182,7 @@ def finalize_session(self):
 
         with summon_full_params_context(self.model, offload_to_cpu=True):
             # in order to update each layer we need to gathers all its parameters
-            finalize()
+            active_session().finalize()
         logger.info("Finalized LLM Compressor session")
         model = get_session_model()
         self.model = model

src/llmcompressor/transformers/finetune/text_generation.py

@@ -49,7 +49,7 @@ def oneshot(**kwargs) -> None:
 
 @deprecated(
     message=(
-        "`from llmcompressor.transformers import train` is deprecated, "
+        "`from llmcompressor import train` is deprecated, "
         "please use `from llmcompressor import train`."
     )
 )

tests/llmcompressor/transformers/finetune/test_finetune_no_recipe_custom_dataset.py

@@ -18,7 +18,7 @@
 
 class TestFinetuneNoRecipeCustomDataset(unittest.TestCase):
     def _test_finetune_wout_recipe_custom_dataset(self):
-        from llmcompressor.transformers import train
+        from llmcompressor import train
 
         dataset_path = Path(tempfile.mkdtemp())
 

tests/llmcompressor/transformers/finetune/test_finetune_without_recipe.py

@@ -20,7 +20,7 @@ def setUp(self):
         self.output = "./finetune_output"
 
     def test_finetune_without_recipe(self):
-        from llmcompressor.transformers import train
+        from llmcompressor import train
 
         recipe_str = None
         device = "cuda:0"

tests/llmcompressor/transformers/finetune/test_oneshot_then_finetune.py

@@ -6,10 +6,9 @@
 from transformers import AutoModelForCausalLM
 from transformers.utils.quantization_config import CompressedTensorsConfig
 
-from llmcompressor import oneshot
+from llmcompressor import oneshot, train
 from llmcompressor.core import create_session
 from llmcompressor.modifiers.quantization import QuantizationModifier
-from llmcompressor.transformers import train
 
 
 @pytest.mark.unit

tests/llmcompressor/transformers/finetune/test_safetensors.py

@@ -22,7 +22,7 @@ def setUp(self):
         self.output = Path("./finetune_output")
 
     def test_safetensors(self):
-        from llmcompressor.transformers import train
+        from llmcompressor import train
 
         device = "cuda:0"
         output_dir = self.output / "output1"

tests/llmcompressor/transformers/test_clear_ml.py

@@ -10,7 +10,7 @@
 except Exception:
     is_clearml = False
 
-from llmcompressor.transformers import train
+from llmcompressor import train
 
 
 @pytest.mark.skipif(not is_clearml, reason="clearML not installed")