[Transforms] Transform Args, Scheme, and Config (#321)

* add args, scheme, config

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

* fix typo

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

* use input/output side

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

* remove llama spinquant

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

* remove side, rename to config_groups

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

* use enum

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

* update docstring, remove unused code

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

---------

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

Author: kylesayrs

src/compressed_tensors/transform/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# flake8: noqa
+# isort: skip_file
+
+from .transform_args import *
+from .transform_scheme import *
+from .transform_config import *

src/compressed_tensors/transform/transform_args.py

@@ -0,0 +1,54 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from enum import Enum
+from typing import Any, List
+
+from pydantic import BaseModel, Field, field_validator
+
+
+__all__ = ["TransformArgs"]
+
+
+class TransformLocation(str, Enum):
+    INPUT = "input"
+    WEIGHT_INPUT = "weight_input"
+    WEIGHT_OUTPUT = "weight_output"
+    OUTPUT = "output"
+    K_CACHE = "k_cache"
+    Q_ATTN = "q_attn"
+
+
+class TransformArgs(BaseModel):
+    """
+    Arguments which define *how* and where a transform should be applied to a model
+
+    :param targets: list of modules to apply transforms to
+    :param location: where to apply transform on module, one of (`input`, `weight`,
+        `output`, `k_cache`, `q_attn`)
+    :param inverse: whether or not to apply the inverse of a transform
+    :param ignore: any modules which should be ignored from the targets list
+    """
+
+    targets: List[str]
+    location: TransformLocation
+    inverse: bool = Field(default=False)
+    ignore: List[str] = Field(default_factory=list)
+
+    @field_validator("targets", "ignore", mode="before")
+    @classmethod
+    def wrap_singleton(cls, value):
+        if isinstance(value, str):
+            return [value]
+        return value

src/compressed_tensors/transform/transform_config.py

@@ -0,0 +1,73 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict
+
+from compressed_tensors.transform import TransformArgs, TransformScheme
+from pydantic import BaseModel
+
+
+__all__ = ["TransformConfig"]
+
+
+class TransformConfig(BaseModel):
+    """
+    Configuration of transforms to be applied to a model. This config is to be
+    serialized within a model's `config.json` file
+
+    :param config_groups: A dictionary of `TransformSchemes` that should be applied
+        to a particular model. The keys can be any arbitrary string
+    """
+
+    config_groups: Dict[str, TransformScheme]
+
+
+# quip / quip sharp
+QUIP = TransformConfig(
+    config_groups={
+        "v": TransformScheme(
+            type="hadamard",
+            apply=[
+                TransformArgs(
+                    targets=["Linear"],
+                    location="input",  # non-mergable
+                ),
+                TransformArgs(
+                    targets=["Linear"],
+                    location="weight_input",
+                    inverse=True,
+                ),
+            ],
+            randomize_modules=True,
+        ),
+        "u": TransformScheme(
+            type="hadamard",
+            apply=[
+                TransformArgs(
+                    targets=["Linear"],
+                    location="weight_output",
+                ),
+                TransformArgs(
+                    targets=["Linear"], location="output", inverse=True  # non-mergable
+                ),
+            ],
+            randomize_modules=True,
+        ),
+    }
+)
+
+
+PRESET_CONFIGS = {
+    "QUIP": QUIP,
+}

src/compressed_tensors/transform/transform_scheme.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List
+
+from compressed_tensors.transform import TransformArgs
+from pydantic import BaseModel, Field
+
+
+__all__ = ["TransformScheme"]
+
+
+class TransformScheme(BaseModel):
+    """
+    Scheme used to parameterize a particular transform type and specify how and where it
+    should be applied to the model
+
+    :param type: string indicating the particular transform type that should be created
+        and applied. This should be one of the registered transform types
+        (see `Transforms.registered_names()`)
+    :param apply: list of TransformationArgs containing the information about the
+        modules that should be targeted by the specified transform
+    :param randomize_modules: True if unique transforms should be applied to each
+        unique module targeted by `apply`, otherwise reuse transform weights where
+        applicable
+    :param requires_grad: True if weights include gradients for training
+    """
+
+    type: str
+    apply: List[TransformArgs] = Field(default_factory=list)
+    randomize_modules: bool = Field(default=False)
+    requires_grad: bool = Field(default=False)

tests/test_transform/test_transform_args.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from compressed_tensors.transform import TransformArgs
+
+
+def test_basic():
+    targets = ["Embedding"]
+    location = "input"
+    args = TransformArgs(targets=targets, location=location)
+
+    assert args.targets == targets
+    assert args.location == location
+    assert len(args.ignore) == 0
+
+
+def test_args_full():
+    targets = ["Linear"]
+    location = "weight_input"
+    inverse = True
+    ignore = ["model.layers.2"]
+
+    args = TransformArgs(
+        targets=targets,
+        location=location,
+        inverse=inverse,
+        ignore=ignore,
+    )
+
+    args.targets = targets
+    args.location == location
+    args.inverse == inverse
+    args.ignore == ignore
+
+
+def test_singleton_targets():
+    target = "target"
+    location = "input"
+    ignore = "ignore"
+    args = TransformArgs(targets=target, location=location, ignore=ignore)
+
+    assert args.targets == [target]
+    assert args.location == location
+    assert args.ignore == [ignore]

tests/test_transform/test_transform_config.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytest
+from compressed_tensors.transform import TransformArgs, TransformConfig, TransformScheme
+
+
+@pytest.fixture
+def basic_transform_scheme():
+    targets = ["Embedding"]
+    location = "input"
+    basic_args = TransformArgs(targets=targets, location=location)
+
+    return TransformScheme(
+        type="hadamard",
+        apply=[basic_args],
+    )
+
+
+def test_basic(basic_transform_scheme):
+    config = TransformConfig(
+        config_groups={
+            "transform_0": basic_transform_scheme,
+        }
+    )
+    assert isinstance(config.config_groups.get("transform_0"), TransformScheme)
+
+
+def test_to_dict(basic_transform_scheme):
+    config = TransformConfig(
+        config_groups={
+            "transform_0": basic_transform_scheme,
+        }
+    )
+    config_dict = config.model_dump()
+    assert "config_groups" in config_dict.keys()
+
+
+def test_multiple_groups():
+    location = "weight_input"
+
+    targets_1 = ["model.layers.0.attn.v_proj"]
+    linear_args_1 = TransformArgs(targets=targets_1, location=location)
+
+    targets_2 = ["model.layers.0.attn.q_proj"]
+    linear_args_2 = TransformArgs(targets=targets_2, location=location)
+
+    scheme_1 = TransformScheme(
+        type="hadamard",
+        apply=[linear_args_1],
+    )
+
+    scheme_2 = TransformScheme(
+        type="hadamard",
+        apply=[linear_args_2],
+    )
+    config = TransformConfig(
+        config_groups={"transform_0": scheme_1, "transform_1": scheme_2}
+    )

tests/test_transform/test_transform_scheme.py

@@ -0,0 +1,74 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from compressed_tensors.transform import TransformArgs, TransformScheme
+
+
+def test_basic_scheme():
+    targets = ["Linear"]
+    location = "input"
+    basic_args = TransformArgs(targets=targets, location=location)
+
+    scheme = TransformScheme(
+        type="hadamard",
+        apply=[basic_args],
+    )
+    assert not scheme.randomize_modules
+    assert scheme.type == "hadamard"
+    assert len(scheme.apply) == 1
+    assert isinstance(scheme.apply[0], TransformArgs)
+
+
+def test_multiple_groups_global():
+    targets = ["Embedding"]
+    location = "input"
+    embedding_args = TransformArgs(targets=targets, location=location)
+
+    targets = ["Linear"]
+    location = "weight_input"
+    linear_args = TransformArgs(targets=targets, location=location)
+
+    # same transform applied to multiple groups
+    scheme = TransformScheme(
+        type="hadamard",
+        apply=[embedding_args, linear_args],
+        randomize_modules=True,
+    )
+
+    assert scheme.randomize_modules
+    assert scheme.type == "hadamard"
+    assert len(scheme.apply) == 2
+    assert isinstance(scheme.apply[0], TransformArgs)
+    assert isinstance(scheme.apply[1], TransformArgs)
+
+
+def test_multiple_groups():
+    apply = []
+    location = "weight_output"
+
+    for i in range(20):
+        targets = [f"model.layers.{i}.attn.v_proj", f"model.layers.{i}.attn.o_proj"]
+        args = TransformArgs(targets=targets, location=location)
+        apply.append(args)
+
+    # global is False, different hadamard transform applied to each group
+    # same dimension/hidden dim
+    scheme = TransformScheme(
+        type="hadamard",
+        apply=apply,
+    )
+
+    assert not scheme.randomize_modules
+    assert scheme.type == "hadamard"
+    assert len(scheme.apply) == 20