[Transforms] Transform Args, Scheme, and Config

[kylesayrs]

Summary

• Introduce Transform Args, Scheme, Config and Data as structures to define recipes to apply transforms
• When compared to [Transforms] Transform Arg, Scheme, and Data support #275, the requirement for ModuleTarget, and TransformData has been removed as well as the transform_creation_args and call_args fields
• Supports spinquant-like transform configs which target non-linear modules such as q_attn and k_cache

TransformArgs

• Arguments which define how and where a transform should be applied to a model

TransformScheme

• Scheme used to parameterize a particular transform type and specify how and where it should be applied to the model
• Includes a list of TransformationArgs onto which the transformation will be applied

TransformConfig

• Configuration of transforms to be applied to a model. This config is to be serialized within a model's config.json file
• The keys can be any arbitrary string and a TransformationScheme should be provided for each new transform type.
• Include preset configs for QUIP/QUIP# and Llama-SpinQuant

Example:

from compressed_tensors.transform import TransformArgs, TransformScheme, TransformsConfig

forward_args = TransformArgs(targets=["Linear"], location="input")  # non-mergable
inverse_args = TransformArgs(targets=["Linear"], location="weight_input", inverse=True)

scheme = TransformScheme(
    type="hadamard",
    apply=[forward_args, inverse_args],
    randomize_modules=True
)

config = TransformsConfig(transform_groups={"v_transform": scheme}

# [transform locations] and {quant locations}:
# input  <- {input @ [input]}
# weight <- {[weight_output] @ weight @ [weight_input]}
# output <- {input @ weight.T @ [output]}

[brian-dellabetta]

Looking good! a few comments

[dsikka]

Two questions:

1. How do we maintain the order of the transforms if multiple are being applied to a given input or weight?
2. What utility will be inferring the transform size and if its right/left?

[kylesayrs]

@dsikka

1. EDIT: order is determined by the order that the schemes are applied in, but the order of application doesn't really matter. The matrix multiplications will still cancel out so long as inputs are prepended
2. FYI @anmarques suggested replacing side ∈ {"left", "right"} with side ∈ {"input", "output"} to clarify how the weight is applied. This is especially nice for quip, since you can directly see the U and V shapes align in the config

To answer the question, this function uses location and side to determine which of the weight to apply to:
https://github.com/neuralmagic/compressed-tensors/pull/316/files#diff-d70cb4816d5b1e1ff0e5f5ea117da8d2e1e106b7419680ba41228a5064dad890R32

[dsikka]

this looks great, thank you for improving on the original design
great work!

[brian-dellabetta]

👍