Accelerate Utilities

[kylesayrs]

Purpose

• Implement offloading utility functions which greatly simplify/clarify offloading-related code in llm-compressor
• Explicitly initialize quantization parameters as offloaded if the module is offloaded

Prerequisites

• Skip accelerate tests #208
• [Bugfix] Update expected shape for per token strategy #210
• add offload_device argument #228

Changes

Changes not covered by prerequisites:

• Implement getattr_chain utility function (also used by llm-compressor)
• Implement depreciated utility decorator for future depreciations
• Implement register_offload_parameter and delete_offload_parameter for easier initialization and removal of parameters related to quantization
• Begin newly initialized quantization parameters on cpu if the module is offloaded offload
  • Faster performance, removes dependency on get_execution_device

Depreciation Strategy

These functions should be depreciated, each for their own reason. These strategies will be implemented in follow-up PRs

Function	Depreciation Reason	Depreciation Strategy
is_module_offloaded	Use official has_offloaded_params	redirect to has_offloaded_params & depreciation warning
get_execution_device	Not useful as a general util	Remove uses from LC & depreciation warning
get_offloaded_device	Folded into update_offload_parameter	Replace uses in LC with update_offload_parameter & depreciation warning
update_prefix_dict	Folded into update_offload_parameter	Replace uses in LC with update_offload_parameter & depreciation warning
update_parameter_data	Use update_offload_data for better args ordering. Open to keeping this one around	Remove uses from LC and CT & depreciation warning

Upstream Strategy

Upstreaming functions to accelerate is a low priority, but comes with the benefit of more reviews and more official support

Function	Upstream Version
register_offload_parameter	N/A
update_offload_data	N/A
delete_offload_parameter	N/A
has_offloaded_params	1.1.0
align_module_device	1.1.0

[rahul-tuli]

LGTM! with a few nits, good work on this!

[dsikka]

What would be the replacement for get_execution_device?

[kylesayrs]

@dsikka The function update_parameter_data takes new_param_data as an input and uses this to update parameter data. Previously, this function would simply overwrite the old data with new_param_data. Now, in order to reduce complexity and increase performance, update_parameter_data requires that parameter being updated and the new parameter data are the same shape.

This assumption causes an error in mock_per_token_calibration, tests/test_quantization/test_configs/test_strategies.py, which revealed to me that the shape used to initialize the per_token strategy and the shape computed by calculate_qparams are different shapes. I consider this ambiguity to be a bug which was causing the test to fail.

[dsikka]

This looks good overall.

Do you mind adding a simple lifecycle dosctring which shows the steps of offloaded modules/parameters to make it slightly easier to follow how the parameters are updated?

I also think we should kick-off W4A16/W8A8 oneshot workflows, similar to what we did here: https://app.asana.com/0/1207078450218847/1208568399648361/f to make sure it runs to completion. I think past issues we've seen have been with g_idx and activation quantization parameters.

[dsikka]

What would be the replacement for get_execution_device?

I think I understand from your PR as to why this can be removed.

[kylesayrs]

@dsikka w.r.t. get_execution_device

1. The function isn't guaranteed to be performant for all device maps, for example half-offloaded models
2. The function has very few uses, so it may be worth removing

For these reasons it's a candidate (and we'll need it for the immediate future), but future work can determine whether we want to keep/ update it