Skip to content

Commit

Permalink
fix fwd func call (#845)
Browse files Browse the repository at this point in the history
Signed-off-by: Kyle Sayers <[email protected]>
  • Loading branch information
dsikka authored and kylesayrs committed Oct 23, 2024
1 parent c6acef2 commit e5e0d9a
Showing 1 changed file with 3 additions and 5 deletions.
8 changes: 3 additions & 5 deletions src/llmcompressor/modifiers/obcq/utils/sgpt_wrapper.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
import time

from compressed_tensors.quantization.lifecycle.forward import (
maybe_calibrate_or_quantize,
)
from compressed_tensors.quantization.lifecycle.forward import forward_quantize

from llmcompressor.modifiers.utils.compression_wrapper import ModuleCompressionWrapper
from llmcompressor.utils import getattr_chain
Expand Down Expand Up @@ -95,7 +93,7 @@ def compress(
args_loc = "quantization_scheme.weights"
weight_quant_args = getattr_chain(self.layer, args_loc, None)
if weight_quant_args is not None:
W = maybe_calibrate_or_quantize(self.layer, W, "weight", weight_quant_args)
W = forward_quantize(self.layer, W, "weight", weight_quant_args)

if isinstance(self.layer, nn.Conv2d):
W = W.flatten(1)
Expand Down Expand Up @@ -213,7 +211,7 @@ def compress(
W = W.t()
W = W.reshape(final_shape).to(final_dtype)
if weight_quant_args is not None:
W = maybe_calibrate_or_quantize(self.layer, W, "weight", weight_quant_args)
W = forward_quantize(self.layer, W, "weight", weight_quant_args)

# This is a bit hacky, but FSDP updates only work if we change the weight in
# place, clone() or direct assignment won't work
Expand Down

0 comments on commit e5e0d9a

Please sign in to comment.