Skip to content
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.

Sign up for GitHub

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

Jump to bottom

[Bug Fix] Fix pruning for partially quantized models #1792

Merged
merged 1 commit into from
Oct 27, 2023
Merged

[Bug Fix] Fix pruning for partially quantized models #1792

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
39 changes: 3 additions & 36 deletions src/sparseml/modifiers/obcq/utils/layer_compressor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,11 @@
from typing import Dict, List

import torch
import torch.nn as nn
from torch.nn import Module

from sparseml.modifiers.obcq.utils.sparsegpt import SparseGPT
from sparseml.pytorch.utils.helpers import get_dependency_order
from sparseml.utils.pytorch.module import get_prunable_layers


_LOGGER = logging.getLogger(__name__)
Expand Down Expand Up @@ -60,14 +60,8 @@ def compressible_modules(self) -> Dict:

:return: dictionary of compressible modules
"""
quantize = self.args.get("quantize", False)
if quantize:
# The layer names are changed due to quantization modifiers, therefore
# we need a slightly different func to retrieve layers
modules = _find_quant_layers(self.layer)
else:
modules = _find_layers(self.layer)
return modules
compressible_layers = get_prunable_layers(self.layer)
return compressible_layers

def pre_compress_parallel(self, **kwargs) -> Dict:
"""
Expand Down Expand Up @@ -217,30 +211,3 @@ def tmp(_, inp, out):
blocksize=self.args["blocksize"],
)
gpts.free()


def _find_quant_layers(
module, layers=[torch.nn.qat.Conv2d, torch.nn.qat.Linear], name=""
):
res = {}
# search for QAT versions of layers
for name1, child in module.named_children():
res.update(
_find_layers(
child, layers=layers, name=name + "." + name1 if name != "" else name1
)
)
return res


def _find_layers(module, layers=[nn.Conv2d, nn.Linear], name=""):
if type(module) in layers:
return {name: module}
res = {}
for name1, child in module.named_children():
res.update(
_find_layers(
child, layers=layers, name=name + "." + name1 if name != "" else name1
)
)
return res
Loading