[Transforms] Apply, serialize, deserialize

[dsikka]

Summary

• Add support to apply transforms to models during quantization, specifically targeting layer weights for now

Process

• Includes processing the provided transforms config, generating a transform_data object to attach to each layer indicating runtime args, and attaching each of the individual transforms to the model layers.
• Specifically, as multiple transforms can be applied to a particular parameter, to differentiate between them, an index value is attached to the parameter name, resulting in transform parameters having the following name convention:{parameter_type}_transform_{idx} e.g. "weight_transform_0" or "input_activation_transform_0" to try and match the convention of weights, input_activations, and output_activations.
• TransformData includes a dictionary of all the transforms-relevant runtime data, and is attached to the layer as "transform_data" . The keys of the dictionary correspond to the transform parameter names. Note: in the future, if we decide to add another layer to infer runtime/call args on the fly, we can potentially remove TransformData but that is an optimization we can talk about in the future.

Apply

• Utils have also been added to apply the transforms to the weights, when applying QDQ. This functionality will be further extended and likely removed from within the forward method as support for activation transforms is added . This is currently being handled by apply_transforms_to_parameter and apply_inverse_transforms_to_parameter which sandwich QDQ

Serialize/Deserialize

• Serialization currently does not compress the transforms and saves them to disk uncompressed (we will either fuse these in or compress them in a follow-up). The quantization_config is also extended with a transforms_config in config.json
• For deserialization, an additional load_transforms function has been added to load the parameters from disk and add the relevant runtime information. However, this requires the above transformers PR indicated above

Examples:

# Apply a transform config to a model

from compressed_tensors.quantization import process_transforms_config

targets = ["Linear"]
module_targets = [ModuleTarget.WEIGHT]
linear_layer_args = TransformationArgs(
    targets=targets, module_targets=module_targets
)

scheme = TransformationScheme(
    transform_type="hadamard",
    groups=[linear_layer_args],
    transform_creation_args={"size": 512},
)
config = TransformationConfig(
    transform_groups={
        "transform_0": scheme,
    }
)

model = torch.nn.Linear(512, 512)

model = process_transforms_config(model=model, transforms_config=config)

# Once processed, the model will have the following parameters:
>> model.weight_transform_0
Parameter containing:
        tensor([[ 1.,  1.,  1.,  ...,  1.,  1.,  1.],
                [ 1., -1.,  1.,  ..., -1.,  1., -1.],
                [ 1.,  1., -1.,  ...,  1., -1., -1.],
                ...,
                [ 1., -1.,  1.,  ..., -1.,  1., -1.],
                [ 1.,  1., -1.,  ...,  1., -1., -1.],
                [ 1., -1., -1.,  ..., -1., -1.,  1.]], dtype=torch.bfloat16)

>> model.transform_data
TransformData(data={'weight_transform_0': 
      { 
          'call_args': defaultdict(),
          'type': "hadamard"
      }
  }

# Apply the transform. 
model(some_dummy_data)

Testing

• This functionality (along with [Transforms] Transform Arg, Scheme, and Data support #275 and [Transforms] Transform Registry Support #274) has been tested using llm-compressor, to quantize and evaluate a Llama 3.1-1b model with the QuantizationModifier, showing performance close to or surpassing GPTQ on gsm8k

Dependencies:

• Requires: [WIP] Add support to load models with transforms huggingface/transformers#36621 to be full functional but should be able to land before then

[brian-dellabetta]

Another example where I think consistency in naming convention will help users. Something like this?

Suggested change

	transforms_config: TransformationConfig,
	transform_config: TransformConfig,

[kylesayrs]

TransformData includes a dictionary of all the transforms-relevant runtime data, and is attached to the layer as "transform_data"

To me, it seems like this information is essentially duplicating the information of the scheme/args? Why not attach the scheme/args to the module, rather than creating a new abstraction?

[brian-dellabetta]

I missed that this is actually modifying model registering transforms to it. The name process_transforms_config is pretty innocuous for something like this, given it's modifying model significantly. Consider renaming to add_transforms_to_model or something more explicit?