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

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!

[kylesayrs]

Fixed a bug, added some tests

[dsikka]

why do we need the per token fix to land as a prereq for this PR?

[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

[dsikka]

this supports non-offloaded modules? for what case?

[kylesayrs]

Supporting non-offloaded modules allows this function to be used throughout the codebase without having to duplicate code

Ugleh...

if not has_offloaded_params(module):
    param = getattr(module, name)
    data = data.to(param.dtype)
    if param.device != "meta":
        param.data.copy_(data)

else:
    update_offload_parameter(module, name, data)

Preetay!

update_offload_parameter(module, name, data)

[dsikka]

When registering the parameters during initialization, don't we know the device already, depending on if the module has been offloaded or not?

We can't pass that device to update_offload_parameter to be used when updating the weights_map?

[kylesayrs]

I've just now rewritten these parts to a bit clearer.

don't we know the device already, depending on if the module has been offloaded or not?

During initialization, the _initialize_scale_zero_point function determines the initial onload device

# begin on the same device as other parameters or cpu if offloaded.
# in the offloaded case, there's no point moving tensors to the execution device
# if they're going to be immediately offloaded by `register_offload_parameter`
params_device = next(module.parameters()).device
device = "cpu" if has_offloaded_params(module) else params_device

It's the job of register_offload_parameter (and by extension update_offload_parameter, offload_to_weights_map) to determine the offload device.

if isinstance(weights_map, dict):
  if key in weights_map:
      offload_device = weights_map[key].device
  else:
      tens = next(iter(weights_map.values()), None)
      offload_device = tens.device if tens is not None else default_device