Skip to content

[Transform] Serialize with tied weights #370

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 92 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
92 commits
Select commit Hold shift + click to select a range
d8a10ec
add utilities
kylesayrs May 30, 2025
d2af054
add tests
kylesayrs May 30, 2025
e32d5b5
add additional tests
kylesayrs May 30, 2025
9d0518b
add utils and tests
kylesayrs May 30, 2025
8c5a2d9
Implement transform factories
kylesayrs May 30, 2025
809e367
Merge branch 'kylesayrs/transform_utils' into kylesayrs/transform_fac…
kylesayrs May 30, 2025
8d613b3
add permutations
kylesayrs May 31, 2025
57d171a
add delete_offload_module
kylesayrs May 31, 2025
d77bcef
Merge branch 'kylesayrs/transform-accelerate-utilities' into kylesayr…
kylesayrs May 31, 2025
ab73b43
Merge branch 'kylesayrs/transform-accelerate-utilities' into kylesayr…
kylesayrs May 31, 2025
4b55733
Merge branch 'kylesayrs/transform_factory' into kylesayrs/transform_p…
kylesayrs May 31, 2025
aa7d21b
key inverses by weight
kylesayrs May 31, 2025
6901e02
fix tests
kylesayrs May 31, 2025
47ae9fe
standardize random hadamard
kylesayrs May 31, 2025
34f1343
Merge branch 'kylesayrs/transform_utils' into kylesayrs/transform_fac…
kylesayrs May 31, 2025
1039100
prepend input hooks
kylesayrs May 31, 2025
5677553
Merge remote-tracking branch 'origin' into kylesayrs/transform_utils
kylesayrs Jun 5, 2025
68ec14e
apply sqrt division first
kylesayrs Jun 5, 2025
a62418a
Merge branch 'kylesayrs/transform_utils' into kylesayrs/transform_fac…
kylesayrs Jun 5, 2025
b117523
use divided hadamards
kylesayrs Jun 5, 2025
a46f754
fix typo
kylesayrs Jun 5, 2025
cb1cb52
add random option
kylesayrs Jun 5, 2025
7c02bb2
Merge branch 'kylesayrs/transform_utils' into kylesayrs/transform_fac…
kylesayrs Jun 5, 2025
02af1e9
use random seeds, rename matrix multiply
kylesayrs Jun 5, 2025
f45f3e9
add deterministic generation to random matrix
kylesayrs Jun 5, 2025
7a7abdf
fix perm math
kylesayrs Jun 5, 2025
6e52894
update docstrings
kylesayrs Jun 5, 2025
7230933
update docstrings
kylesayrs Jun 5, 2025
f74fe3e
Merge branch 'kylesayrs/transform_factory' into kylesayrs/transform_p…
kylesayrs Jun 5, 2025
92ddea9
cleanup
kylesayrs Jun 5, 2025
779956f
cleanup 2
kylesayrs Jun 5, 2025
fbd2939
Merge branch 'kylesayrs/transform_utils' into kylesayrs/transform_fac…
kylesayrs Jun 5, 2025
dd72b6a
make seed optional
kylesayrs Jun 5, 2025
4ae491d
Merge branch 'kylesayrs/transform_factory' into kylesayrs/transform_p…
kylesayrs Jun 5, 2025
da19b0f
remove iterable check and missing return value
kylesayrs Jun 9, 2025
7ab17ce
Merge branch 'main' into kylesayrs/transform_permutations
kylesayrs Jun 10, 2025
33df50f
Merge remote-tracking branch 'origin' into kylesayrs/transform_permut…
kylesayrs Jun 10, 2025
6e1ec39
Remove unrelated changes
kylesayrs Jun 10, 2025
938e702
simplify code
kylesayrs Jun 10, 2025
27bc0b3
implement apply, use in tests
kylesayrs Jun 10, 2025
a27db62
use hadamards database file
kylesayrs Jun 11, 2025
ce63955
try manifest
kylesayrs Jun 11, 2025
7ae5863
try setup, update hadamards list
kylesayrs Jun 11, 2025
67675c3
fix setup
kylesayrs Jun 11, 2025
f061db9
add docstrings, cleanup
kylesayrs Jun 11, 2025
4a84ce1
fix setup, thank you @dbarbuzzi
kylesayrs Jun 11, 2025
cde1066
remove numpy, add tests
kylesayrs Jun 11, 2025
1ba6195
solidify dtype, add gpu tests
kylesayrs Jun 11, 2025
c373345
fix docstring
kylesayrs Jun 11, 2025
fbaf47a
add device option
kylesayrs Jun 11, 2025
5a887f4
construct on execution device, cache on offload device
kylesayrs Jun 11, 2025
310fe6d
save construction device changes for later
kylesayrs Jun 11, 2025
b715329
construct on execution device, cache on offload device
kylesayrs Jun 11, 2025
249323c
cite nja sloane
kylesayrs Jun 11, 2025
1823af4
Merge branch 'kylesayrs/extend-hadamard', remote-tracking branch 'ori…
kylesayrs Jun 11, 2025
94a0bf5
Merge remote-tracking branch 'origin' into kylesayrs/extend-hadamard
kylesayrs Jun 11, 2025
cf066e0
Merge branch 'kylesayrs/extend-hadamard' into kylesayrs/transform_con…
kylesayrs Jun 11, 2025
c1a4a34
remove dreg
kylesayrs Jun 11, 2025
5807ee1
put on device via safe_open
kylesayrs Jun 11, 2025
ccb88ed
nits and docstrings
kylesayrs Jun 12, 2025
feba695
update docstring
kylesayrs Jun 12, 2025
c8f6b53
Merge branch 'kylesayrs/extend-hadamard' into kylesayrs/transform_con…
kylesayrs Jun 12, 2025
e7f08e1
Merge branch 'kylesayrs/transform_construct_cache_device' into kylesa…
kylesayrs Jun 12, 2025
75b9307
Merge remote-tracking branch 'origin' into kylesayrs/transform_permut…
kylesayrs Jun 12, 2025
b6a0dd4
Merge remote-tracking branch 'origin' into kylesayrs/transform_constr…
kylesayrs Jun 13, 2025
955f2f5
Merge
kylesayrs Jun 23, 2025
226f367
merge with construct: construct in float32
kylesayrs Jun 23, 2025
9745acb
Merge remote-tracking branch 'origin' into kylesayrs/transform_apply
kylesayrs Jun 23, 2025
fd3390a
construct with same dtype, constructing on fp32 found no difference
kylesayrs Jun 23, 2025
3c55003
Merge branch 'kylesayrs/transform_construct_cache_device' into kylesa…
kylesayrs Jun 23, 2025
ad29c15
remove unnecessary imports
kylesayrs Jun 23, 2025
85f40b5
bugfixes (#375)
brian-dellabetta Jul 2, 2025
500af9b
use factory_kwargs
kylesayrs Jul 7, 2025
8e36540
add frozen dict to deps
kylesayrs Jul 7, 2025
48653ec
Merge remote-tracking branch 'origin' into kylesayrs/transform_permut…
kylesayrs Jul 7, 2025
56df0f7
fix style
kylesayrs Jul 7, 2025
a251569
merge
kylesayrs Jul 7, 2025
cb5a32b
Merge remote-tracking branch 'origin' into kylesayrs/transform_apply
kylesayrs Jul 7, 2025
06e0346
Merge branch 'kylesayrs/transform_permutations' into kylesayrs/transf…
kylesayrs Jul 7, 2025
0a4fea5
Merge branch 'kylesayrs/transform_construct_cache_device' into kylesa…
kylesayrs Jul 7, 2025
49740c6
use delete_offload_module
kylesayrs Jul 7, 2025
7dc182b
Merge remote-tracking branch 'origin' into kylesayrs/transform_constr…
kylesayrs Jul 7, 2025
80db2ce
Merge branch 'kylesayrs/transform_construct_cache_device' into kylesa…
kylesayrs Jul 7, 2025
e06bbad
add docstrign
kylesayrs Jul 7, 2025
438bc13
Merge remote-tracking branch 'origin' into kylesayrs/transform_apply
kylesayrs Jul 7, 2025
fd77ecc
use parametrize
kylesayrs Jul 8, 2025
5a95fd2
populate _dynamic_tied_weights_keys
kylesayrs Jun 28, 2025
b009f47
ensure serializable
kylesayrs Jul 8, 2025
2e362d2
remove extra space
kylesayrs Jul 8, 2025
c6abb96
apply style
kylesayrs Jul 8, 2025
3da59a0
Merge remote-tracking branch 'origin' into kylesayrs/transform_save
kylesayrs Jul 10, 2025
97345b0
merge dregs
kylesayrs Jul 10, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 51 additions & 3 deletions src/compressed_tensors/transform/factory/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,8 @@
# limitations under the License.

from abc import ABC, abstractmethod
from typing import Optional
from collections import defaultdict
from typing import List, Optional, Tuple

import torch
import torch.nn.utils.parametrize as P
Expand Down Expand Up @@ -49,10 +50,13 @@ class TransformFactory(RegistryMixin, ABC):
:param seed: random seed used to transform weight randomization
"""

transforms: List["TransformBase"]

def __init__(self, name: str, scheme: TransformScheme, seed: Optional[int] = None):
self.name = name
self.scheme = scheme
self.generator = torch.Generator()
self.transforms = list()
if seed is not None:
self.generator.manual_seed(seed)

Expand Down Expand Up @@ -91,16 +95,26 @@ def apply_to_model(self, model: Module):
if is_target(name, module, arg.targets, arg.ignore):
self._apply_to_module(module, arg)

self._update_tied_weights()

def _apply_to_module(self, module: Module, args: TransformArgs):
"""
Create transforms and apply them to the module

:param module: target module to apply transforms to
:param args: defines how the transform will be applied to the target module
"""
if has_offloaded_params(module):
if module._hf_hook.place_submodules:
raise NotImplementedError(
"Applying transforms to offloaded submodules with "
"`place_submodules=True` is not supported"
)

# create transform as submodule
transform_name = f"{self.name}_{args.location.value}"
transform = self.create_transform(module, args)
self.transforms.append(transform)
register_offload_module(module, transform_name, transform)

# register input transformation hook
Expand Down Expand Up @@ -131,8 +145,9 @@ def input_hook(_, args):
raise ValueError("Offloaded training is not supported")
P.register_parametrization(module, "weight", transform)

# transform is no longer needed (unfusing is not supported)
delete_offload_module(module, transform_name)
else:
# transform is no longer needed (unfusing is not supported)
delete_offload_module(module, transform_name)

# register output transformation hook
elif args.location == TransformLocation.OUTPUT:
Expand All @@ -146,6 +161,34 @@ def output_hook(_, _input, output):
else:
raise NotImplementedError()

def _update_tied_weights(self):
"""
Populate the `_dynamic_tied_weights_keys` attribute of transforms,
which is used by transformers to detect and remove shared pointers
during saving
"""
# avoid issues with this method being called twice
for transform in self.transforms:
transform._dynamic_tied_weights_keys = list()

# map from data_ptrs to keys
ptr_to_keys: dict[int, List[Tuple[TransformBase, str]]] = defaultdict(list)
for transform in self.transforms:
for name, param in transform.named_parameters(recurse=False):
# NOTE: previously asserted that parent._hf_hook.place_submodules=False
if has_offloaded_params(transform):
param = transform._hf_hook.weights_map[name]
ptr_to_keys[param.data_ptr()].append((transform, name))

# populate `_dynamic_tied_weights_keys` if there is more than one key
# and ensure that they share tensors
for shared_keys in ptr_to_keys.values():
if len(shared_keys) > 1:
tensor = getattr(shared_keys[0][0], shared_keys[0][1])

for transform, name in shared_keys:
transform._dynamic_tied_weights_keys.append(name)
setattr(transform, name, tensor)

class TransformBase(InternalModule, ABC):
"""
Expand All @@ -154,6 +197,11 @@ class TransformBase(InternalModule, ABC):

args: TransformArgs
weight: Parameter
_dynamic_tied_weights_keys: List[str]
Copy link
Contributor

@brian-dellabetta brian-dellabetta Jul 8, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this be a set instead to avoid this issue commented above?
# avoid issues with this method being called twice

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm a little confused how setting this would avoid issues with double-calling _update_tied_weights?


def __init__(self):
super().__init__()
self._dynamic_tied_weights_keys = list()

@abstractmethod
def forward(self, value: Tensor) -> Tensor:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ def _create_weight(self, size: int, dtype: dtype, device: device) -> Parameter:

def _create_inverse(self, weight: Parameter) -> Parameter:
data = high_precision_invert(weight.data)
data = data.contiguous() # ensure proper serialization
return Parameter(data, requires_grad=False)


Expand Down
7 changes: 4 additions & 3 deletions tests/test_transform/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,13 @@

import pytest
import torch
from compressed_tensors.transform import TransformArgs
from compressed_tensors.transform import TransformArgs, TransformFactory
from transformers import PretrainedConfig, PreTrainedModel


class TransformableModel(torch.nn.Module):
class TransformableModel(PreTrainedModel):
def __init__(self, *sizes):
super().__init__()
super().__init__(config=PretrainedConfig())
self.fcs = torch.nn.ModuleList(
[
torch.nn.Linear(sizes[index], sizes[index + 1], bias=False)
Expand Down
50 changes: 50 additions & 0 deletions tests/test_transform/factory/test_serialization.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
# 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
import torch
from compressed_tensors.transform import (
TransformConfig,
TransformScheme,
apply_transform_config,
)
from compressed_tensors.utils import offloaded_dispatch
from tests.test_transform.conftest import scheme_kwargs
from tests.testing_utils import requires_accelerate, requires_gpu


@pytest.mark.parametrize("scheme_kwargs", scheme_kwargs())
def test_serialization(scheme_kwargs, model_apply, tmp_path, offload=False):
# get model, maybe offload
model, apply = model_apply
if offload:
offloaded_dispatch(model, torch.device("cuda"))

# apply transforms to model
config = TransformConfig(
config_groups={"": TransformScheme(**scheme_kwargs, apply=apply)}
)
apply_transform_config(model, config)

# save model
model.save_pretrained(tmp_path)

# TODO: reload model


@requires_gpu
@requires_accelerate()
@pytest.mark.parametrize("scheme_kwargs", scheme_kwargs())
def test_serialization_offload(scheme_kwargs, model_apply, tmp_path):
test_serialization(scheme_kwargs, model_apply, tmp_path, offload=True)
Loading