transform arg support

Author: dsikka

src/compressed_tensors/transforms/transform_args.py

@@ -0,0 +1,68 @@
+# 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 collections import defaultdict
+from enum import Enum
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Field, field_validator
+
+
+__all__ = ["TransformationArgs", "ModuleTarget"]
+
+# TODO: we eventually want to target generic parameters but for now, this
+# is sufficient
+class ModuleTarget(str, Enum):
+    """
+    Enum storing parameter or activation types being targeted by transforms
+    in a particuilar module.
+    """
+
+    WEIGHT = "weight"
+    INPUT_ACTIVATIONS = "input_activations"
+    OUTPUT_ACTIVATIONS = "output_activations"
+
+    @classmethod
+    def has_value(cls, value):
+        return value in cls._value2member_map_
+
+
+class TransformationArgs(BaseModel):
+    """
+    User-facing arguments used to define which modules and their specific
+    parameters and/or activations should be targeted by a particular transform.
+
+    :param targets: list of layers to target
+    :param module_targets: list of layer parameters and/or activations onto which the
+        transform should be applied. The same transform will be applied for all
+        module targets in the list.
+    :param call_args: dictionary of args needed for the transform during runtime,
+        beyond the input_tensor or transform
+    :param ignore: any submodule which should be ignored from the targets list
+
+    """
+
+    targets: List[str]
+    module_targets: List[Union[ModuleTarget, str]]
+    call_args: Optional[Dict[str, Any]] = defaultdict()
+    ignore: Optional[List[str]] = Field(default_factory=list)
+
+    @field_validator("module_targets", mode="before")
+    def validate_module_target(cls, value) -> List[ModuleTarget]:
+        module_targets_list = []
+        for v in value:
+            assert ModuleTarget.has_value(v.lower())
+            module_targets_list.append(v)
+
+        return module_targets_list

src/compressed_tensors/transforms/transform_config.py

@@ -0,0 +1,37 @@
+# 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.transforms.transform_scheme import TransformationScheme
+from pydantic import BaseModel
+
+
+__all__ = ["TransformationConfig"]
+
+
+class TransformationConfig(BaseModel):
+    """
+    Configuration of transforms to be added within a model's config.json.
+
+    :param transform_groups: A dictionary of the different TransformationSchemes
+        that should be applied to a particular model. The keys can be any
+        arbitrary string and a TransformationScheme should be provided
+        for each new transform type.
+    """
+
+    transform_groups: Dict[str, TransformationScheme]
+
+    def to_dict(self):
+        return self.model_dump()

src/compressed_tensors/transforms/transform_data.py

@@ -0,0 +1,35 @@
+# 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 dataclasses import dataclass
+from typing import Any, Callable, Dict
+
+
+__all__ = ["TransformData"]
+
+
+# TODO: adding for now but we may not need it during runtime depending on the
+# integration.
+@dataclass
+class TransformData:
+    """
+    Data that is required during runtime in order to apply the transform.
+
+    Example:
+        data={transform_name: {"apply": Callable, "call_args": Dict}})
+        transform_data = TransformData(data=data)
+    """
+
+    data: Dict
+    idx: int = 0

src/compressed_tensors/transforms/transform_scheme.py

@@ -0,0 +1,48 @@
+# 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 Any, Dict, List, Optional
+
+from compressed_tensors.transforms import Transforms
+from compressed_tensors.transforms.transform_args import TransformationArgs
+from pydantic import BaseModel, field_validator
+
+
+__all__ = ["TransformationScheme"]
+
+
+class TransformationScheme(BaseModel):
+    """
+    :param transform_type: string indicating the particular transform that
+        should be created and applied. This should be one of the registered transforms
+        i.e be in Transforms.registered_names()
+    :param groups: includes TransformationArgs containing the information about the
+        layers that should be targeted by the specified transform. By providing a list,
+        users have the ability to apply the same transform type (with the same set
+        of arguments) to different layers.
+    :param transform_creation_args: arguments needed to initialize the transform, if
+        any
+    :param global_transform: whether an identical transform is applied to all the
+        TransformationArgs in the groups list
+    """
+
+    transform_type: str
+    groups: List[TransformationArgs]
+    global_transform: bool = False
+    transform_creation_args: Optional[Dict[str, Any]] = None
+
+    @field_validator("transform_type", mode="before")
+    def validate_transform_type(cls, value) -> str:
+        assert value in Transforms.registered_names()
+        return value

tests/test_transforms/test_transform_args.py

@@ -0,0 +1,52 @@
+# 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 collections import defaultdict
+
+from compressed_tensors.transforms.transform_args import (
+    ModuleTarget,
+    TransformationArgs,
+)
+
+
+def test_transform_args_basic():
+    targets = ["Embedding"]
+    module_targets = [ModuleTarget.INPUT_ACTIVATIONS]
+    basic_args = TransformationArgs(targets=targets, module_targets=module_targets)
+
+    assert basic_args.targets[0] == "Embedding"
+    assert basic_args.module_targets[0] == ModuleTarget.INPUT_ACTIVATIONS
+    assert isinstance(type(basic_args.call_args), type(defaultdict))
+    assert len(basic_args.ignore) == 0
+
+
+def test_transform_args_full():
+    targets = ["Linear"]
+    module_targets = ["weight", "input_activations"]
+    ignore = ["model.layers.2"]
+    call_args = {"transpose": True}
+
+    full_args = TransformationArgs(
+        targets=targets,
+        module_targets=module_targets,
+        call_args=call_args,
+        ignore=ignore,
+    )
+
+    full_args.targets = targets
+    full_args.ignore == ignore
+    full_args.module_targets[0] == ModuleTarget.WEIGHT.value
+    full_args.module_targets[1] == ModuleTarget.INPUT_ACTIVATIONS.value
+    assert full_args.call_args.get("transpose")

tests/test_transforms/test_transform_config.py

@@ -0,0 +1,80 @@
+# 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.transforms.transform_args import (
+    ModuleTarget,
+    TransformationArgs,
+)
+from compressed_tensors.transforms.transform_config import TransformationConfig
+from compressed_tensors.transforms.transform_scheme import TransformationScheme
+
+
+@pytest.fixture
+def basic_transform_scheme():
+    targets = ["Embedding"]
+    module_targets = [ModuleTarget.INPUT_ACTIVATIONS]
+    basic_args = TransformationArgs(targets=targets, module_targets=module_targets)
+
+    scheme = TransformationScheme(
+        transform_type="hadamard",
+        groups=[basic_args],
+        transform_creation_args={"size": 1024},
+    )
+    return scheme
+
+
+def test_basic(basic_transform_scheme):
+    config = TransformationConfig(
+        transform_groups={
+            "transform_0": basic_transform_scheme,
+        }
+    )
+    assert isinstance(config.transform_groups.get("transform_0"), TransformationScheme)
+
+
+def test_to_dict(basic_transform_scheme):
+    config = TransformationConfig(
+        transform_groups={
+            "transform_0": basic_transform_scheme,
+        }
+    )
+    config_dict = config.to_dict()
+    assert "transform_groups" in config_dict.keys()
+
+
+def test_multiple_groups():
+    module_targets = [ModuleTarget.WEIGHT]
+
+    targets_1 = ["model.layers.0.attn.v_proj"]
+    linear_args_1 = TransformationArgs(targets=targets_1, module_targets=module_targets)
+
+    targets_2 = ["model.layers.0.attn.q_proj"]
+    linear_args_2 = TransformationArgs(targets=targets_2, module_targets=module_targets)
+
+    scheme_1 = TransformationScheme(
+        transform_type="hadamard",
+        groups=[linear_args_1],
+        transform_creation_args={"size": 1024},
+    )
+
+    scheme_2 = TransformationScheme(
+        transform_type="hadamard",
+        groups=[linear_args_2],
+        transform_creation_args={"size": 256},
+    )
+    config = TransformationConfig(
+        transform_groups={"transform_0": scheme_1, "transform_1": scheme_2}
+    )