Skip to content
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.

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

[Export][Transformers] Enable loading text-generation datasets #1938

Merged

Conversation

dbogunowicz
Copy link
Contributor

Feature description

From now on, whenever we export the text-generation model, we are also able to load the appropriate dataset.
This is achieved by integrating @Satrat latest registry for LLM dataloaders into the transformers export.

Testing

export(
            source_path=source_path,
            target_path=target_path,
            task='text-generation',
            num_export_samples=4,
            **dict(
                data_args=dict(
                    dataset_name="wikitext", dataset_config_name="wikitext-2-raw-v1"
                )
            ),
        )
Fetching 10 files: 100%|██████████| 10/10 [00:00<00:00, 52.54it/s]
...
2024-01-04 18:00:59 sparseml.export.export INFO     Exporting 4 samples...
4it [00:02,  1.61it/s]
2024-01-04 18:01:01 sparseml.export.export_data INFO     Exporting sample-inputs to /tmp/pytest-of-damian/pytest-0/test_export_samples_roneneldan0/target...
2024-01-04 18:01:01 sparseml.export.export_data INFO     Successfully exported sample-inputs to /tmp/pytest-of-damian/pytest-0/test_export_samples_roneneldan0/target!
2024-01-04 18:01:01 sparseml.export.export_data INFO     Exporting sample-outputs to /tmp/pytest-of-damian/pytest-0/test_export_samples_roneneldan0/target...
2024-01-04 18:02:46 sparseml.export.export_data INFO     Successfully exported sample-outputs to /tmp/pytest-of-damian/pytest-0/test_export_samples_roneneldan0/target!
...
2024-01-04 18:02:47 sparseml.exporters.transforms.kv_cache.transforms_codegen INFO     Successfully adjusted the causal_mask input
2024-01-04 18:02:47 sparseml.exporters.transforms.onnx_transform INFO     [AdditionalTransformsCodeGen] Transformed 10 matches
2024-01-04 18:02:47 sparseml.export.helpers INFO     Optimization: kv_cache_injection has been successfully applied to the ONNX model: /tmp/pytest-of-damian/pytest-0/test_export_samples_roneneldan0/target/deployment/model.onnx
2024-01-04 18:02:47 sparseml.export.export INFO     Validating model structure...
2024-01-04 18:02:47 sparseml.export.validators WARNING  File /tmp/pytest-of-damian/pytest-0/test_export_samples_roneneldan0/target/sample-labels is missing.
2024-01-04 18:02:47 sparseml.export.validators WARNING  File /tmp/pytest-of-damian/pytest-0/test_export_samples_roneneldan0/target/deployment/tokenizer.model is missing.
2024-01-04 18:02:47 sparseml.export.export INFO     Successfully exported model from:
/tmp/pytest-of-damian/pytest-0/test_export_samples_roneneldan0/target
to
/tmp/pytest-of-damian/pytest-0/test_export_samples_roneneldan0/target/deployment
for integration: transformers

:return: The output of the model with the past key values unpacked
"""

past_key_values = output_vals["past_key_values"]
Copy link
Contributor

Choose a reason for hiding this comment

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

nit: can we make this a constant rather than hard coding it? Since its used in a few places

Copy link
Contributor Author

Choose a reason for hiding this comment

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

This is removed now all together

@@ -196,6 +198,28 @@ def run_inference_with_dict_data(
return inputs, labels, output


def _unnest_past_key_values(output_vals: Dict[str, Any]) -> Dict[str, Any]:
Copy link
Contributor

Choose a reason for hiding this comment

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

Not sure I'm following what this is for, can you give a quick example of what output_vals looks like before and after this function? Why is it that output_vals contains these past_key_values in the first place?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

This is removed now all together

* 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
@dbogunowicz dbogunowicz merged commit 6179cb2 into feature/damian/feature_branch_export Jan 5, 2024
@dbogunowicz dbogunowicz deleted the feature/damian/samples_llms branch January 5, 2024 20:15
bfineran added a commit that referenced this pull request Jan 10, 2024
* 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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants