[Export Refactor] Prepare the module to be more general (before including `transformers`) #1908

dbogunowicz · 2023-12-15T08:51:41Z

No description provided.

src/sparseml/export/export.py

src/sparseml/export/export_data.py

src/sparseml/pytorch/image_classification/integration_helper_functions.py

…unctions.py

…s.py

src/sparseml/export/export.py

bfineran · 2023-12-15T19:59:31Z

src/sparseml/export/export.py

+    _LOGGER.info(
+        f"Applying optimizations: {graph_optimizations} to the exported model..."
+    )
+    apply_optimizations(


what is the purpose of apply_optimizations top level function over helper_functions.graph_optimizations again worried about specificity to onnx

graph_optimization contains a mapping that specifies all the possible optimizations for a specific integration. the apply_optimization is a general function that establishes which of those optimizations should be actually applied and applies them. Why worried about onnx specificity? We could always pass the deployment_target as an argument of apply_optimizations - this actually demonstrates the usefulness of this abstraction.

I still fail to see how folding this function into each of the IntegrationHelperFunctions is beneficial. It would add to code duplication, and IMO worsen the readability for the user. I'd propose to leave it for now, it's a very small component, can be refactored easily if needed.

src/sparseml/export/export_data.py

Co-authored-by: Benjamin Fineran <bfineran@users.noreply.github.com>

* cleanup * Delete src/sparseml/transformers/integration_helper_functions_generative.py * Delete src/sparseml/transformers/utils/optimizations.py * Delete tests/sparseml/export/transformers/test_generative_transformers.py * Delete tests/sparseml/transformers/test_integration_helper_functions_generative.py * addressing PR reviews * [Export Refactor] Export generative transformers(#1910)

…ndent of the task name

…sparseml into feature/damian/export_adapt

dbogunowicz · 2023-12-19T10:13:39Z

src/sparseml/export/export.py

@@ -112,6 +117,7 @@ def export(

    # choose the appropriate device
    device = default_device() if device == "auto" else device
+    device = use_single_gpu(device) if "cuda" in device else device


We want to be in full control over the placement of the model on the device.
User can either specify the device:str ourselves (to cpu, cuda, cuda:0) or keep the default auto configuration. If auto, we use default_device() helper function that returns a proper device string.
Question 1: shouldn't we default to cpu here?
Question 2: if requested, should we allow the user to place the model over multiple CUDA devices?

Added the use_single_gpu function

for general initiailizaitaon -> model to multiple gpus modelparallel

src/sparseml/export/export.py

dbogunowicz · 2023-12-19T10:14:37Z

src/sparseml/export/export.py

@@ -126,69 +132,55 @@ def export(
    _LOGGER.info(f"Starting export for {integration} model...")

    helper_functions: IntegrationHelperFunctions = (
-        IntegrationHelperFunctions.load_from_registry(integration)
+        IntegrationHelperFunctions.load_from_registry(integration, task=task)


Moved the use of the optional task argument here.

dbogunowicz · 2023-12-19T10:15:08Z

src/sparseml/export/export.py

+    # loaded_model_kwargs may include any objects
+    # that were created along with the model and are needed
+    # for the export
+    model, loaded_model_kwargs = helper_functions.create_model(


auxiliary_items renamed to loaded_model_args

dbogunowicz · 2023-12-19T10:18:38Z

src/sparseml/export/export_data.py

+    for batch_num, data in tqdm(enumerate(data_loader)):
+        if batch_num == num_samples:
+            break
+        if isinstance(data, dict):


as requested before, we can run inference on the transformer data but now without the mandatory export from src.transformers.

dbogunowicz · 2023-12-19T10:19:19Z

src/sparseml/export/helpers.py

+
+
+def create_export_kwargs(
+    loaded_model_kwargs: Dict[str, Any], export_target: str = "deepsparse"


added export_target argument for the future, in case we need different logic for different export types (moving beyond onnx)

dbogunowicz · 2023-12-19T10:20:06Z

src/sparseml/integration_helper_functions.py

@@ -32,10 +32,12 @@ class Integrations(Enum):
    """

    image_classification = "image-classification"
+    transformers = "transformers"


removed transformers_generative from available options

dbogunowicz · 2023-12-19T10:21:14Z

src/sparseml/pytorch/utils/helpers.py

@@ -137,6 +137,15 @@ def default_device() -> str:
    return "cuda:{}".format(",".join(device_ids))


+def use_single_gpu(device: str) -> str:


Potentially useful, related to the open question in to the comment left in exporters.py

dbogunowicz · 2023-12-19T10:21:40Z

src/sparseml/pytorch/utils/helpers.py

@@ -538,7 +547,8 @@ def _tensors_export_batch(
        return

    if isinstance(tensors, Iterable):
-        for index, tens in enumerate(zip(*tensors)):
+        # TODO: I am breaking something here? - dbogunowicz
+        for index, tens in enumerate(zip(tensors)):


(tests will tell once we attempt to land to main)

dbogunowicz · 2023-12-19T10:22:34Z

src/sparseml/transformers/integration_helper_functions.py

+        task = kwargs.get("task")
+        if task is None:
+            raise ValueError("To create a transformer model, a task must be specified")
+        if task in TaskNames.text_generation.value:


This is what we discussed Ben - now the transformers IntegrationHelperFunctions can dynamically change their fields given the task argument.

dbogunowicz · 2023-12-19T10:23:32Z

src/sparseml/transformers/utils/initializers.py

+        recipe_args=None,
+        teacher=None,
+    )
+    applied = trainer.apply_manager(epoch=math.inf, checkpoint=None)


Will be syncing with Sara today regarding the new recipy application

dbogunowicz · 2023-12-19T11:56:42Z

src/sparseml/transformers/integration_helper_functions.py

+        elif task in TaskNames.text_generation.value:
+            # if the task is text generation, alter the default attributes
+            # to reflect the idiosyncrasies for text generation
+            self.apply_optimizations = apply_optimizations_generative_transformer


folded the apply_optimizations function into the IntegrationHelperFunctions

dbogunowicz · 2023-12-19T15:35:33Z

Merging to feature branch per @bfineran approval to accelerate working on important features.

…ding `transformers`) (#1908) * adapt the export script to handle transformers * Update src/sparseml/pytorch/image_classification/integration_helper_functions.py * Delete tests/sparseml/export/transformers/__init__.py * Delete tests/sparseml/export/transformers/test_generative_transformers.py * Delete tests/sparseml/export/transformers/test_transformers.py * Update src/sparseml/export/export.py Co-authored-by: Benjamin Fineran <bfineran@users.noreply.github.com> * addressing review comments * [Export Refactor] Export `transformers` (#1909) * cleanup * Delete src/sparseml/transformers/integration_helper_functions_generative.py * Delete src/sparseml/transformers/utils/optimizations.py * Delete tests/sparseml/export/transformers/test_generative_transformers.py * Delete tests/sparseml/transformers/test_integration_helper_functions_generative.py * addressing PR reviews * [Export Refactor] Export generative transformers(#1910) * make tests green, remove using task to resolve the integration type * fix all the tests after the merge, make integration resolution independent of the task name * fold generative transformers into transformer helper functions * complete tests for export_data.py * Update src/sparseml/export/export.py * add tests that confirms that kv cache injection has been added * move applying optimizations into integration helper functions --------- Co-authored-by: Benjamin Fineran <bfineran@users.noreply.github.com>

* initial commit * respond to PR comments * [Export Refactor][Image Classification] `create_model` function (#1878) * initial commit * looking good, time to cleanup * Delete src/sparseml/export/helpers.py * Delete tests/sparseml/export/test_helpers.py * ready for review * improve design * tests pass * reuse _validate_dataset_num_classes * [Export Refactor][Image Classification] `create_dummy_input` function (#1880) * initial commit * looking good, time to cleanup * Delete src/sparseml/export/helpers.py * Delete tests/sparseml/export/test_helpers.py * ready for review * improve design * tests pass * reuse _validate_dataset_num_classes * initial commit * Update src/sparseml/pytorch/image_classification/integration_helper_functions.py * Update src/sparseml/pytorch/image_classification/integration_helper_functions.py * ready for review * Update src/sparseml/export/export.py * Update src/sparseml/integration_helper_functions.py * [Export Refactor][Image Classification] `export_model` function (#1883) * initial commit * looking good, time to cleanup * Delete src/sparseml/export/helpers.py * Delete tests/sparseml/export/test_helpers.py * ready for review * improve design * tests pass * reuse _validate_dataset_num_classes * initial commit * Update src/sparseml/pytorch/image_classification/integration_helper_functions.py * Update src/sparseml/pytorch/image_classification/integration_helper_functions.py * ready for review * Update src/sparseml/export/export.py * Update src/sparseml/integration_helper_functions.py * initial commit * fixes * ready for review * nit * add return * make export function more general * [Export Refactor][Image Classification] `apply_optimizations` function (#1884) * initial commit * looking good, time to cleanup * Delete src/sparseml/export/helpers.py * Delete tests/sparseml/export/test_helpers.py * ready for review * improve design * tests pass * reuse _validate_dataset_num_classes * initial commit * Update src/sparseml/pytorch/image_classification/integration_helper_functions.py * Update src/sparseml/pytorch/image_classification/integration_helper_functions.py * ready for review * Update src/sparseml/export/export.py * Update src/sparseml/integration_helper_functions.py * initial commit * fixes * ready for review * nit * add return * initial commit * [Export Refactor][Image Classification] `export_sample_inputs_outputs` function (#1888) * initial commit * looking good, time to cleanup * Delete src/sparseml/export/helpers.py * Delete tests/sparseml/export/test_helpers.py * ready for review * improve design * tests pass * reuse _validate_dataset_num_classes * initial commit * Update src/sparseml/pytorch/image_classification/integration_helper_functions.py * Update src/sparseml/pytorch/image_classification/integration_helper_functions.py * ready for review * Update src/sparseml/export/export.py * Update src/sparseml/integration_helper_functions.py * initial commit * fixes * ready for review * nit * add return * initial commit * initial commit * PR comments * beautification * remove duplicated function * [Export Refactor][Image Classification] `create_deployment_folder` function (#1889) * initial commit * looking good, time to cleanup * Delete src/sparseml/export/helpers.py * Delete tests/sparseml/export/test_helpers.py * ready for review * improve design * tests pass * reuse _validate_dataset_num_classes * initial commit * Update src/sparseml/pytorch/image_classification/integration_helper_functions.py * Update src/sparseml/pytorch/image_classification/integration_helper_functions.py * ready for review * Update src/sparseml/export/export.py * Update src/sparseml/integration_helper_functions.py * initial commit * fixes * ready for review * nit * add return * initial commit * initial commit * initial commit * fix rebase, tests_work * ready to push * [Export Refactor][Image Classification] `validate_correctness` function (#1890) * initial commit * looking good, time to cleanup * Delete src/sparseml/export/helpers.py * Delete tests/sparseml/export/test_helpers.py * ready for review * improve design * tests pass * reuse _validate_dataset_num_classes * initial commit * Update src/sparseml/pytorch/image_classification/integration_helper_functions.py * Update src/sparseml/pytorch/image_classification/integration_helper_functions.py * ready for review * Update src/sparseml/export/export.py * Update src/sparseml/integration_helper_functions.py * initial commit * fixes * ready for review * nit * add return * initial commit * initial commit * initial commit * initial commit * Delete tests/sparseml/test_integration_helper_functions.py * ready to merge * [Export Refactor] End to end testing (#1898) * initial commit * looking good, time to cleanup * Delete src/sparseml/export/helpers.py * Delete tests/sparseml/export/test_helpers.py * ready for review * improve design * tests pass * reuse _validate_dataset_num_classes * initial commit * Update src/sparseml/pytorch/image_classification/integration_helper_functions.py * Update src/sparseml/pytorch/image_classification/integration_helper_functions.py * ready for review * Update src/sparseml/export/export.py * Update src/sparseml/integration_helper_functions.py * initial commit * fixes * ready for review * nit * add return * initial commit * initial commit * initial commit * initial commit * Delete tests/sparseml/test_integration_helper_functions.py * ready to merge * add structure validator * ready for review * Delete tests/sparseml/export/model.onnx * Delete tests/sparseml/export/image_classification/model.onnx * Delete tests/sparseml/export/image_classification/conftest.py * PR comments * remove onnx * [Export Refactor] Prepare the module to be more general (before including `transformers`) (#1908) * adapt the export script to handle transformers * Update src/sparseml/pytorch/image_classification/integration_helper_functions.py * Delete tests/sparseml/export/transformers/__init__.py * Delete tests/sparseml/export/transformers/test_generative_transformers.py * Delete tests/sparseml/export/transformers/test_transformers.py * Update src/sparseml/export/export.py Co-authored-by: Benjamin Fineran <bfineran@users.noreply.github.com> * addressing review comments * [Export Refactor] Export `transformers` (#1909) * cleanup * Delete src/sparseml/transformers/integration_helper_functions_generative.py * Delete src/sparseml/transformers/utils/optimizations.py * Delete tests/sparseml/export/transformers/test_generative_transformers.py * Delete tests/sparseml/transformers/test_integration_helper_functions_generative.py * addressing PR reviews * [Export Refactor] Export generative transformers(#1910) * make tests green, remove using task to resolve the integration type * fix all the tests after the merge, make integration resolution independent of the task name * fold generative transformers into transformer helper functions * complete tests for export_data.py * Update src/sparseml/export/export.py * add tests that confirms that kv cache injection has been added * move applying optimizations into integration helper functions --------- Co-authored-by: Benjamin Fineran <bfineran@users.noreply.github.com> * [Export Refactor][Transformers] Enable loading SparseModels (#1921) * initial commit * adressing review comments * Fix the tests * fix tests with help from sara * [Export][Transformers] Enable loading `text-generation` datasets (#1938) * add suport for past_key_values in sample-outputs * [Export][Transformers] Implementation of correctness validation (#1935) * fix tests with help from sara * Update src/sparseml/transformers/utils/initializers.py * swap sparsezoo validator for custom one (top k match) * add more informative error message * add correctness validation for LLMs * remove past_key_values from outputs * remove past_key_values from outputs (2) * small note comment for the future * tests fixed * fix test * [Export refactor] final manual testing fixes (#1948) * [Export refactor] final manual testing fixes * review --------- Co-authored-by: Benjamin Fineran <bfineran@users.noreply.github.com>

adapt the export script to handle transformers

23e892f

dbogunowicz requested a review from bfineran December 15, 2023 08:51

dbogunowicz commented Dec 15, 2023

View reviewed changes

src/sparseml/export/export.py Outdated Show resolved Hide resolved

dbogunowicz commented Dec 15, 2023

View reviewed changes

src/sparseml/export/export_data.py Outdated Show resolved Hide resolved

dbogunowicz commented Dec 15, 2023

View reviewed changes

src/sparseml/export/export_data.py Show resolved Hide resolved

dbogunowicz commented Dec 15, 2023

View reviewed changes

src/sparseml/pytorch/image_classification/integration_helper_functions.py Outdated Show resolved Hide resolved

dbogunowicz added 4 commits December 15, 2023 09:58

Update src/sparseml/pytorch/image_classification/integration_helper_f…

2bb1f20

…unctions.py

Delete tests/sparseml/export/transformers/__init__.py

fb20ae2

Delete tests/sparseml/export/transformers/test_generative_transformer…

eab9fb1

…s.py

Delete tests/sparseml/export/transformers/test_transformers.py

0897615

dbogunowicz mentioned this pull request Dec 15, 2023

[Export Refactor] Feature Branch #1858

Merged

19 tasks

bfineran reviewed Dec 15, 2023

View reviewed changes

dbogunowicz and others added 2 commits December 18, 2023 08:37

Update src/sparseml/export/export.py

9c7fac5

Co-authored-by: Benjamin Fineran <bfineran@users.noreply.github.com>

addressing review comments

05a9ee3

dbogunowicz requested a review from bfineran December 18, 2023 08:11

dbogunowicz and others added 6 commits December 18, 2023 16:40

make tests green, remove using task to resolve the integration type

313ecf8

fix all the tests after the merge, make integration resolution indepe…

6829514

…ndent of the task name

Merge branch 'feature/damian/export_adapt' of github.com:neuralmagic/…

4ada09a

…sparseml into feature/damian/export_adapt

fold generative transformers into transformer helper functions

5840eb5

complete tests for export_data.py

d7c2375

dbogunowicz commented Dec 19, 2023

View reviewed changes

src/sparseml/export/export.py Outdated Show resolved Hide resolved

Update src/sparseml/export/export.py

c308949

dbogunowicz commented Dec 19, 2023

View reviewed changes

dbogunowicz added 2 commits December 19, 2023 10:35

add tests that confirms that kv cache injection has been added

ad2a7f0

move applying optimizations into integration helper functions

9fbead1

dbogunowicz commented Dec 19, 2023

View reviewed changes

dbogunowicz merged commit 506d7fe into feature/damian/feature_branch_export Dec 19, 2023

dbogunowicz deleted the feature/damian/export_adapt branch December 19, 2023 15:35

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Export Refactor] Prepare the module to be more general (before including `transformers`) #1908

[Export Refactor] Prepare the module to be more general (before including `transformers`) #1908

dbogunowicz commented Dec 15, 2023

bfineran Dec 15, 2023

dbogunowicz Dec 18, 2023

dbogunowicz Dec 19, 2023

dbogunowicz Dec 19, 2023

dbogunowicz Dec 19, 2023

dbogunowicz Dec 19, 2023

dbogunowicz Dec 19, 2023

dbogunowicz Dec 19, 2023

dbogunowicz Dec 19, 2023

dbogunowicz Dec 19, 2023

dbogunowicz Dec 19, 2023

dbogunowicz Dec 19, 2023

dbogunowicz Dec 19, 2023

dbogunowicz Dec 19, 2023

dbogunowicz Dec 19, 2023

dbogunowicz commented Dec 19, 2023



		def create_export_kwargs(
		loaded_model_kwargs: Dict[str, Any], export_target: str = "deepsparse"

		@@ -137,6 +137,15 @@ def default_device() -> str:
		return "cuda:{}".format(",".join(device_ids))


		def use_single_gpu(device: str) -> str:

[Export Refactor] Prepare the module to be more general (before including transformers) #1908

[Export Refactor] Prepare the module to be more general (before including transformers) #1908

Conversation

dbogunowicz commented Dec 15, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dbogunowicz commented Dec 19, 2023

[Export Refactor] Prepare the module to be more general (before including `transformers`) #1908

[Export Refactor] Prepare the module to be more general (before including `transformers`) #1908