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

Accelerate Utilities Follow-up #224

Merged

Conversation

kylesayrs
Copy link
Contributor

@kylesayrs kylesayrs commented Dec 6, 2024

Purpose

  • Add more utilities not covered by Accelerate Utilities #193
  • Add tests for offloading functions
  • Reduce code complexity by using utility functions
  • Fix bug in update_offload_data and expand capabilities

Changes

  • Implement align_module_device, which has already been upstreamed and is just copied here until we bump the required accelerate version
  • Implement disable_hf_hook and utilize in initialize_module_for_quantization
    • Note that the previous "dictionary recreation" logic is no longer needed. It was previously necessary since the newly attached parameters were not updated in the state dict, however this is now done by register_offload_parameter
  • Fix bug in update_offload_data where data would not be updated properly if the module was offloaded and weights_map dataset was being improperly gotten
  • Cover raw dictionary case in update_offload_data

Testing

  • Added tests in tests/test_utils/test_offload.py
  • Added offloaded test in tests/test_quantization/lifecycle/test_initialize.py

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
@kylesayrs kylesayrs merged commit e7e1d81 into kylesayrs/upstream-candidates Dec 6, 2024
@kylesayrs kylesayrs deleted the kylesayrs/upstream-candidates-2 branch December 6, 2024 06:39
dsikka pushed a commit that referenced this pull request Dec 20, 2024
* wip

* add modify_offload_module

* update docs

* WIP

* cleanup functions, begin depreciation

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

* remove extra space

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

* revert get_offloaded_device

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

* update to align_module_device

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

* add requires skip for accelerate

* fix per token initialization

* remove align_module_device

* respond to nits

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

* Accelerate Utilities Follow-up (#224)

* rename

* implement recursive case

* remove print

* support OffloadedWeightsLoader

* add lifecycle docstring

* implement offload_to_weights_map with recursive definition

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

* add docstring

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

* fix type hint

* add check_accelerate guard

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

* make device used by  clearer

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

* update update_prefix_dict

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

* reuse fixture

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

* use apply rather than recursion

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

* clearer delete_from_weights_map

* add offload_device argument (#228)

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

---------

Signed-off-by: Kyle Sayers <kylesayrs@gmail.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.

1 participant