mct-nightly 2.2.0.20241203.546__py3-none-any.whl → 2.2.0.20241205.533__py3-none-any.whl

model_compression_toolkit/gptq/keras/gptq_training.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Callable, List, Tuple, Union
+from typing import Callable, List, Tuple, Union, Generator
 
 import tensorflow as tf
 from keras import Model
@@ -20,11 +20,13 @@ from packaging import version
 from tensorflow.keras.layers import Layer
 from tqdm import tqdm
 
-from model_compression_toolkit.core.common.hessian import HessianInfoService
+from model_compression_toolkit.core.common.hessian import HessianInfoService, HessianScoresGranularity
 # As from Tensorflow 2.6, keras is a separate package and some classes should be imported differently.
 from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.core.keras.back2framework.keras_model_builder import KerasModelBuilder
-from model_compression_toolkit.core.keras.data_util import data_gen_to_dataloader
+from model_compression_toolkit.core.keras.data_util import data_gen_to_dataloader, \
+    FixedSampleInfoDataset, FixedTFDataset, create_tf_dataloader, TFDatasetFromGenerator, \
+    IterableSampleWithConstInfoDataset
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
 from model_compression_toolkit.gptq.common.gradual_activation_quantization import \
     get_gradual_activation_quantizer_wrapper_factory
@@ -83,13 +85,10 @@ class KerasGPTQTrainer(GPTQTrainer):
 
         """
 
-        def _get_total_grad_steps():
-            return get_total_grad_steps(representative_data_gen) * gptq_config.n_epochs
-
-        # This must be set before the model building (as it is required for activation holder construction),
-        # which occurs in the base constructor.
-        self.gradual_act_quantizer_wrapper_factory = get_gradual_activation_quantizer_wrapper_factory(
-            gptq_config, _get_total_grad_steps, KerasLinearAnnealingScheduler)
+        self.fw_soft_quantizer_regularization = SoftQuantizerRegularization
+        self.fw_linear_annealing_scheduler = KerasLinearAnnealingScheduler
+        self.fw_get_gptq_trainable_parameters_fn = get_gptq_trainable_parameters
+        self.fw_get_weights_for_loss_fn = get_weights_for_loss
 
         super().__init__(graph_float,
                          graph_quant,
@@ -99,53 +98,106 @@ class KerasGPTQTrainer(GPTQTrainer):
                          representative_data_gen_fn=representative_data_gen,
                          hessian_info_service=hessian_info_service)
 
-        self.loss_list = []
-        self.input_scale = 1
-
-        trainable_weights, bias_weights, trainable_threshold = get_gptq_trainable_parameters(
-            self.fxp_model,
-            fw_info,
-            add_bias=gptq_config.train_bias)
-
-        self.flp_weights_list, self.fxp_weights_list = get_weights_for_loss(self.fxp_model)
-
-        if not (len(self.compare_points) == len(trainable_weights) == len(self.flp_weights_list) == len(
-                self.fxp_weights_list)):
-            Logger.critical("Mismatch in the number of comparison points, layers with trainable weights, "
-                            "and the number of float and quantized weights for loss calculation. "
-                            "Ensure all these elements align to proceed with GPTQ training.")
-
-        flattened_trainable_weights = [w for layer_weights in trainable_weights for w in layer_weights]
-        flattened_bias_weights = [w for layer_weights in bias_weights for w in layer_weights]
-        trainable_quantization_parameters = trainable_threshold
-        self.optimizer_with_param = self.get_optimizer_with_param(flattened_trainable_weights,
-                                                                  flattened_bias_weights,
-                                                                  trainable_quantization_parameters)
-        self.has_params_to_train = np.sum(
-            [len(optimizer_params_tuple[1]) for optimizer_params_tuple in self.optimizer_with_param]) > 0
-
-        if self.float_user_info.input_scale != self.gptq_user_info.input_scale:
-            Logger.critical("Input scale mismatch detected between the float model and the GPTQ model. "
-                            "Confirm that the input scales for both models are correctly configured and aligned.")  # pragma: no cover
-        else:
-            self.input_scale = self.gptq_user_info.input_scale
 
-        self.weights_for_average_loss = self._get_compare_points_loss_weights()
+    def _prepare_train_dataloader_sla(self, data_gen_fn: Callable[[], Generator]) -> tf.data.Dataset:
+        """
+        Computes Sample-Layer Attention score and builds a train dataloader in TensorFlow.
 
-        self.reg_func = get_regularization(self.gptq_config,
-                                           _get_total_grad_steps,
-                                           SoftQuantizerRegularization,
-                                           KerasLinearAnnealingScheduler)
+        Args:
+            data_gen_fn: function for representative dataset generation.
 
-    def _get_compare_points_loss_weights(self):
-        """ Get compare points weights for the distillation loss. """
-        if self.gptq_config.hessian_weights_config:
-            hess_dataloader = data_gen_to_dataloader(self.representative_data_gen_fn,
-                                                     batch_size=self.gptq_config.hessian_weights_config.hessian_batch_size)
-            return self.compute_hessian_based_weights(hess_dataloader)
+        Returns:
+            TensorFlow dataset yielding three outputs - samples, weights for the distillation loss,
+            and weights for regularization.
+        """
+        # Create a fixed dataset
+        fixed_dataset = FixedTFDataset(data_gen_fn)
+        orig_batch_size = fixed_dataset.orig_batch_size
+
+        # Prepare a separate loader for computing hessians over the whole dataset
+        hess_data_loader = create_tf_dataloader(
+            fixed_dataset,
+            batch_size=self.gptq_config.hessian_weights_config.hessian_batch_size,
+            shuffle=False
+        )
+
+        # Prepare request for Hessian computation
+        request = self._build_hessian_request(
+            granularity=HessianScoresGranularity.PER_OUTPUT_CHANNEL,
+            data_loader=hess_data_loader,
+            n_samples=None
+        )
+        layers_hessians = self.hessian_service.fetch_hessian(request, force_compute=True)
+
+        # Compute SLA score defined as max over elements
+        layers_hessians = {
+            layer: tf.convert_to_tensor(tf.reduce_max(hess, axis=tuple(range(1, len(hess.shape))))) for layer, hess in layers_hessians.items()
+        }
+
+        # Stack hessians for comparison points
+        hessians_tensor = tf.stack([layers_hessians[layer.name] for layer in self.compare_points])
+        assert hessians_tensor.shape[0] == len(self.compare_points)
+        loss_weights = list(hessians_tensor.numpy())  # Convert to a list for compatibility
+
+        # Prepare final dataset with samples and loss weights
+        sla_train_dataset = FixedSampleInfoDataset(fixed_dataset.samples, loss_weights)
+
+        # Calculate regularization weights as mean across samples
+        reg_weights = tf.reduce_mean(hessians_tensor, axis=1)
+
+        # Define a collate function to add regularization weights to each batch
+        def collate_fn(samples_with_loss_weights):
+            return *samples_with_loss_weights, reg_weights
+
+        # Create final dataset using the new dataloader with collate_fn
+        final_dataset = create_tf_dataloader(
+            dataset=sla_train_dataset,
+            batch_size=orig_batch_size,
+            shuffle=True,
+            collate_fn=collate_fn
+        )
+
+        return final_dataset
+
+    def _prepare_train_dataloader_for_non_sla(self,
+                                              data_gen_fn: Callable[[], Generator]) -> tf.data.Dataset:
+        """
+        Prepares a train dataloader for non-SLA tasks.
 
+        Args:
+            data_gen_fn: Factory for representative dataset generator.
+
+        Returns:
+            A `tf.data.Dataset` yielding samples with loss weights and regularization weights.
+        """
+        # Step 1: Create a dataset from the generator
+        dataset = TFDatasetFromGenerator(data_gen_fn)
         num_nodes = len(self.compare_points)
-        return np.ones((num_nodes,)) / num_nodes
+
+        # Step 2: Compute loss weights
+        if self.gptq_config.hessian_weights_config:
+            hessian_dataset = create_tf_dataloader(dataset=dataset, batch_size=self.gptq_config.hessian_weights_config.hessian_batch_size)
+            hessian_weights = self.compute_hessian_based_weights(hessian_dataset)
+            loss_weights = tf.convert_to_tensor(hessian_weights, dtype=tf.float32)
+        else:
+            loss_weights = tf.ones(num_nodes, dtype=tf.float32) / num_nodes
+
+        # Step 3: Create a dataset with samples and loss weights
+        augmented_dataset = IterableSampleWithConstInfoDataset(dataset.dataset, loss_weights)
+
+        # Step 4: Add constant regularization weights
+        reg_weights = tf.ones(num_nodes, dtype=tf.float32)
+
+        def collate_fn(batch):
+            samples, loss_weights = batch
+            return samples, loss_weights, reg_weights
+
+        # Step 5: Create a tf.data.Dataset with collate_fn
+        train_dataloader = create_tf_dataloader(augmented_dataset,
+                                                batch_size=dataset.orig_batch_size,
+                                                collate_fn=collate_fn)
+
+        return train_dataloader
 
     def _is_gptq_weights_trainable(self,
                                    node: common.BaseNode) -> bool:
@@ -226,9 +278,13 @@ class KerasGPTQTrainer(GPTQTrainer):
 
         return gptq_model, gptq_user_info
 
-    def compute_gradients(self, in_y_float: List[tf.Tensor], input_data: List[np.ndarray],
+    def compute_gradients(self,
+                          in_y_float: List[tf.Tensor],
+                          input_data: List[np.ndarray],
                           in_optimizer_with_param: List,
-                          training=True) -> Tuple[tf.Tensor, List[tf.Tensor]]:
+                          training=True,
+                          distill_loss_weights=None,
+                          reg_weights=None) -> Tuple[tf.Tensor, List[tf.Tensor]]:
         """
         Get outputs from both teacher and student networks. Compute the observed error,
         and use it to compute the gradients and applying them to the student weights.
@@ -253,9 +309,9 @@ class KerasGPTQTrainer(GPTQTrainer):
                                                self.flp_weights_list,
                                                self.compare_points_mean,
                                                self.compare_points_std,
-                                               self.weights_for_average_loss)
+                                               distill_loss_weights)
 
-            reg_value = self.reg_func(self.fxp_model, self.gptq_config.regularization_factor)
+            reg_value = self.reg_func(self.fxp_model, self.gptq_config.regularization_factor, reg_weights)
 
             loss_value += reg_value
 
@@ -279,14 +335,19 @@ class KerasGPTQTrainer(GPTQTrainer):
         # Training loop
         # ----------------------------------------------
         if self.has_params_to_train:
-            self.micro_training_loop(self.representative_data_gen_fn,
-                                     compute_gradients,
+            self.micro_training_loop(compute_gradients,
                                      self.optimizer_with_param,
                                      self.gptq_config.n_epochs,
                                      True)
 
     @tf.function
-    def nano_training_step(self, input_data, in_compute_gradients, in_optimizer_with_param, is_training):
+    def nano_training_step(self,
+                           input_data,
+                           in_compute_gradients,
+                           in_optimizer_with_param,
+                           is_training,
+                           distill_loss_weights,
+                           reg_weights):
         """
         This function run part of the training step, wrapped by a tf.function for acceleration.
         Args:
@@ -303,12 +364,15 @@ class KerasGPTQTrainer(GPTQTrainer):
         # run float model
         y_float = self.float_model(input_data)
         # rung quantized model and calculate loss & gradients
-        loss_value_step, grads = in_compute_gradients(y_float, input_data, in_optimizer_with_param,
-                                                      training=is_training)
+        loss_value_step, grads = in_compute_gradients(y_float,
+                                                      input_data,
+                                                      in_optimizer_with_param,
+                                                      training=is_training,
+                                                      distill_loss_weights=distill_loss_weights,
+                                                      reg_weights=reg_weights)
         return loss_value_step, grads
 
     def micro_training_loop(self,
-                            data_function: Callable,
                             in_compute_gradients: Callable,
                             in_optimizer_with_param: List[Tuple[tf.keras.optimizers.Optimizer, List[tf.Tensor]]],
                             n_epochs: int,
@@ -316,7 +380,6 @@ class KerasGPTQTrainer(GPTQTrainer):
         """
         This function run a micro training loop on given set of parameters.
         Args:
-            data_function: A callable function that give a batch of samples.
             in_compute_gradients: A callable function that compute the gradients.
             in_optimizer_with_param: A list of optimizer classes to update with the corresponding parameters.
             n_epochs: Number of update iterations of representative dataset.
@@ -327,12 +390,19 @@ class KerasGPTQTrainer(GPTQTrainer):
         """
         with tqdm(range(n_epochs), "Running GPTQ optimization") as epochs_pbar:
             for _ in epochs_pbar:
-                with tqdm(data_function(), position=1, leave=False) as data_pbar:
+                with tqdm(self.train_dataloader, position=1, leave=False) as data_pbar:
                     for data in data_pbar:
-                        input_data = [d * self.input_scale for d in data]
 
-                        loss_value_step, grads = self.nano_training_step(input_data, in_compute_gradients,
-                                                                         in_optimizer_with_param, is_training)
+                        input_data, distill_loss_weights, reg_weight = data
+
+                        input_data = [d * self.input_scale for d in input_data]
+
+                        loss_value_step, grads = self.nano_training_step(input_data,
+                                                                         in_compute_gradients,
+                                                                         in_optimizer_with_param,
+                                                                         is_training,
+                                                                         distill_loss_weights,
+                                                                         reg_weight)
                         # Run one step of gradient descent by updating
                         # the value of the variables to minimize the loss.
                         for i, (o, p) in enumerate(in_optimizer_with_param):

model_compression_toolkit/gptq/keras/graph_info.py

@@ -16,7 +16,6 @@
 import tensorflow as tf
 from typing import Tuple, List
 from model_compression_toolkit.core.keras.constants import USE_BIAS
-from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from tensorflow.keras.models import Model
 from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
@@ -26,7 +25,6 @@ from model_compression_toolkit.trainable_infrastructure.common.base_trainable_qu
 
 
 def get_gptq_trainable_parameters(fxp_model: Model,
-                                  fw_info: FrameworkInfo,
                                   add_bias: bool = False) -> (
         List[tf.Variable], List[tf.Variable], List[tf.Variable]):
     """
@@ -34,7 +32,6 @@ def get_gptq_trainable_parameters(fxp_model: Model,
 
     Args:
         fxp_model: Model to get its trainable parameters.
-        fw_info: Framework information needed for keras kernel ops list.
         add_bias: Whether to include biases of the model (if there are) or not.
 
     Returns:
@@ -60,7 +57,7 @@ def get_gptq_trainable_parameters(fxp_model: Model,
             trainable_threshold.extend(quantizer_trainable_threshold)
 
             if add_bias:
-                kernel_ops_attrs = fw_info.kernel_ops_attributes_mapping.get(type(layer.layer))
+                kernel_ops_attrs = DEFAULT_KERAS_INFO.kernel_ops_attributes_mapping.get(type(layer.layer))
                 use_bias = kernel_ops_attrs is not None and kernel_ops_attrs[0] is not None \
                            and layer.layer.get_config().get(USE_BIAS)
                 if use_bias is not None and use_bias and layer.layer.bias is not None:

model_compression_toolkit/gptq/keras/quantization_facade.py

@@ -19,7 +19,7 @@ from packaging import version
 
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.gptq.common.gptq_constants import REG_DEFAULT, LR_DEFAULT, LR_REST_DEFAULT, \
-    LR_BIAS_DEFAULT, GPTQ_MOMENTUM
+    LR_BIAS_DEFAULT, GPTQ_MOMENTUM, REG_DEFAULT_SLA
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import TENSORFLOW, ACT_HESSIAN_DEFAULT_BATCH_SIZE, GPTQ_HESSIAN_NUM_SAMPLES
 from model_compression_toolkit.verify_packages import FOUND_TF
@@ -42,7 +42,7 @@ if FOUND_TF:
     from model_compression_toolkit.gptq.keras.gptq_keras_implementation import GPTQKerasImplemantation
     from model_compression_toolkit.core.keras.keras_model_validation import KerasModelValidation
     from tensorflow.keras.models import Model
-    from model_compression_toolkit.gptq.keras.gptq_loss import GPTQMultipleTensorsLoss
+    from model_compression_toolkit.gptq.keras.gptq_loss import GPTQMultipleTensorsLoss, sample_layer_attention_loss
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from model_compression_toolkit.exporter.model_wrapper import get_exportable_keras_model
     from model_compression_toolkit import get_target_platform_capabilities
@@ -61,11 +61,12 @@ if FOUND_TF:
     def get_keras_gptq_config(n_epochs: int,
                               optimizer: OptimizerV2 = None,
                               optimizer_rest: OptimizerV2 = None,
-                              loss: Callable = GPTQMultipleTensorsLoss(),
+                              loss: Callable = None,
                               log_function: Callable = None,
                               use_hessian_based_weights: bool = True,
-                              regularization_factor: float = REG_DEFAULT,
+                              regularization_factor: float = None,
                               hessian_batch_size: int = ACT_HESSIAN_DEFAULT_BATCH_SIZE,
+                              use_hessian_sample_attention: bool = False,
                               gradual_activation_quantization: Union[bool, GradualActivationQuantizationConfig] = False) -> GradientPTQConfig:
         """
         Create a GradientPTQConfig instance for Keras models.
@@ -79,6 +80,7 @@ if FOUND_TF:
             use_hessian_based_weights (bool): Whether to use Hessian-based weights for weighted average loss.
             regularization_factor (float): A floating point number that defines the regularization factor.
             hessian_batch_size (int): Batch size for Hessian computation in Hessian-based weights GPTQ.
+            use_hessian_sample_attention (bool): whether to use Sample-Layer Attention score for weighted loss.
             gradual_activation_quantization (bool, GradualActivationQuantizationConfig): If False, GradualActivationQuantization is disabled. If True, GradualActivationQuantization is enabled with the default settings. GradualActivationQuantizationConfig object can be passed to use non-default settings.
 
         returns:
@@ -105,9 +107,25 @@ if FOUND_TF:
         """
         optimizer = optimizer or tf.keras.optimizers.Adam(learning_rate=LR_DEFAULT)
         optimizer_rest = optimizer_rest or tf.keras.optimizers.Adam(learning_rate=LR_REST_DEFAULT)
+        bias_optimizer = tf.keras.optimizers.SGD(learning_rate=LR_BIAS_DEFAULT, momentum=GPTQ_MOMENTUM)
 
-        bias_optimizer = tf.keras.optimizers.SGD(learning_rate=LR_BIAS_DEFAULT,
-                                                 momentum=GPTQ_MOMENTUM)
+        if regularization_factor is None:
+            regularization_factor = REG_DEFAULT_SLA if use_hessian_sample_attention else REG_DEFAULT
+
+        loss = loss or GPTQMultipleTensorsLoss()
+        hessian_weights_config = None
+        if use_hessian_sample_attention:
+            if not use_hessian_based_weights:    # pragma: no cover
+                raise ValueError('use_hessian_based_weights must be set to True in order to use Sample Layer Attention.')
+
+            hessian_weights_config = GPTQHessianScoresConfig(per_sample=True,
+                                                             hessians_num_samples=None,
+                                                             hessian_batch_size=hessian_batch_size)
+            loss = loss or sample_layer_attention_loss
+        elif use_hessian_based_weights:
+            hessian_weights_config = GPTQHessianScoresConfig(per_sample=False,
+                                                             hessians_num_samples=GPTQ_HESSIAN_NUM_SAMPLES,
+                                                             hessian_batch_size=hessian_batch_size)
 
         if isinstance(gradual_activation_quantization, bool):
             gradual_quant_config = GradualActivationQuantizationConfig() if gradual_activation_quantization else None
@@ -117,11 +135,6 @@ if FOUND_TF:
             raise TypeError(f'gradual_activation_quantization argument should be bool or '
                             f'GradualActivationQuantizationConfig, received {type(gradual_activation_quantization)}')
 
-        hessian_weights_config = None
-        if use_hessian_based_weights:
-            hessian_weights_config = GPTQHessianScoresConfig(per_sample=False,
-                                                             hessians_num_samples=GPTQ_HESSIAN_NUM_SAMPLES,
-                                                             hessian_batch_size=hessian_batch_size)
         return GradientPTQConfig(n_epochs=n_epochs,
                                  optimizer=optimizer,
                                  optimizer_rest=optimizer_rest,

model_compression_toolkit/gptq/keras/quantizer/soft_rounding/soft_quantizer_reg.py

@@ -40,30 +40,42 @@ class SoftQuantizerRegularization:
         self.count_iter = tf.Variable(0.)
 
 
-    def __call__(self, model: Model, entropy_reg: float):
+    def __call__(self, model: Model, entropy_reg: float, layer_weights: tf.Tensor):
         """
         Returns the soft quantizer regularization value for SoftRounding.
 
         Args:
             model: A model to be quantized with SoftRounding.
             entropy_reg: Entropy value to scale the quantizer regularization.
+            layer_weights: a vector of layers weights.
 
         Returns: Regularization value.
         """
-        soft_reg_aux: List[tf.Tensor] = []
+        layers = [l for l in model.layers if isinstance(l, KerasTrainableQuantizationWrapper)]
+
+        if layer_weights.shape[0] != len(layers):
+            raise ValueError(f'Expected weights.shape[0] to be {len(layers)}, '
+                             f'received shape {layer_weights.shape}.')  # pragma: no cover
+
         b = self.beta_scheduler(self.count_iter.value())
-        for layer in model.layers:
-            if isinstance(layer, KerasTrainableQuantizationWrapper):
-                kernel_attribute = get_kernel_attribute_name_for_gptq(layer_type=type(layer.layer),
-                                                                      fw_info=DEFAULT_KERAS_INFO)
 
-                st = layer.weights_quantizers[kernel_attribute].get_soft_targets()
-                soft_reg_aux.append(tf.reduce_sum(1 - tf.pow(tf.math.abs(st - .5) * 2, b)))
+        max_w = tf.reduce_max(layer_weights)
+
+        # Initialize reg to zero
+        reg = tf.constant(0.0, dtype=tf.float32)
+
+        # Compute the regularization term without concatenating
+        for i, layer in enumerate(layers):
+            kernel_attribute = get_kernel_attribute_name_for_gptq(layer_type=type(layer.layer),
+                                                                  fw_info=DEFAULT_KERAS_INFO)
+
+            st = layer.weights_quantizers[kernel_attribute].get_soft_targets()
 
-        reg = 0
+            soft_loss = tf.reduce_sum(1 - tf.pow(tf.math.abs(st - 0.5) * 2, b))
+            reg += layer_weights[i] * soft_loss
 
-        for sq in soft_reg_aux:
-            reg += sq
+        # Normalize reg by max_w
+        reg = reg / max_w
 
         self.count_iter.assign_add(1.0)
 

model_compression_toolkit/gptq/pytorch/gptq_training.py

@@ -21,9 +21,6 @@ from torch.nn import Module
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
-from model_compression_toolkit.gptq.common.gradual_activation_quantization import get_gradual_activation_quantizer_wrapper_factory
-from model_compression_toolkit.gptq.common.regularization_factory import get_regularization
-
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
@@ -41,7 +38,6 @@ from model_compression_toolkit.gptq.pytorch.graph_info import get_gptq_trainable
 from model_compression_toolkit.gptq.pytorch.quantizer.quantization_builder import quantization_builder
 
 from mct_quantizers import PytorchQuantizationWrapper, PytorchActivationQuantizationHolder
-from model_compression_toolkit.trainable_infrastructure.common.util import get_total_grad_steps
 from model_compression_toolkit.trainable_infrastructure.pytorch.annealing_schedulers import PytorchLinearAnnealingScheduler
 from model_compression_toolkit.gptq.pytorch.quantizer.soft_rounding.soft_quantizer_reg import SoftQuantizerRegularization as PytorchSoftQuantizerRegularization
 
@@ -76,13 +72,10 @@ class PytorchGPTQTrainer(GPTQTrainer):
             representative_data_gen: Dataset to use for inputs of the models.
             hessian_info_service: HessianInfoService to fetch info based on the hessian approximation of the float model.
         """
-        def _get_total_grad_steps():
-            # TODO get it from the dataset
-            return get_total_grad_steps(representative_data_gen) * gptq_config.n_epochs
-
-        # must be set prior to model building in the base class constructor
-        self.gradual_act_quantizer_wrapper_factory = get_gradual_activation_quantizer_wrapper_factory(
-            gptq_config, _get_total_grad_steps, PytorchLinearAnnealingScheduler)
+        self.fw_soft_quantizer_regularization = PytorchSoftQuantizerRegularization
+        self.fw_linear_annealing_scheduler = PytorchLinearAnnealingScheduler
+        self.fw_get_gptq_trainable_parameters_fn = get_gptq_trainable_parameters
+        self.fw_get_weights_for_loss_fn = get_weights_for_loss
 
         super().__init__(graph_float,
                          graph_quant,
@@ -92,40 +85,6 @@ class PytorchGPTQTrainer(GPTQTrainer):
                          representative_data_gen_fn=representative_data_gen,
                          hessian_info_service=hessian_info_service)
 
-        self.loss_list = []
-        self.input_scale = 1
-        if self.float_user_info.input_scale != self.gptq_user_info.input_scale:
-            Logger.critical("Input scale mismatch between float and GPTQ networks. "
-                            "Ensure both networks have matching input scales.")  # pragma: no cover
-        else:
-            self.input_scale = self.gptq_user_info.input_scale
-
-        trainable_weights, trainable_bias, trainable_threshold = get_gptq_trainable_parameters(
-            self.fxp_model,
-            add_bias=self.gptq_config.train_bias)
-
-        self.flp_weights_list, self.fxp_weights_list = get_weights_for_loss(self.fxp_model)
-        if not (len(self.compare_points) == len(trainable_weights) == len(self.flp_weights_list) == len(
-                self.fxp_weights_list)):
-            Logger.critical("GPTQ: Number of comparison points, layers with trainable weights, "
-                            "and float vs. quantized weights for loss calculation do not match. "
-                            "Verify consistency across these parameters for successful GPTQ training.")
-
-        self.optimizer_with_param = self.get_optimizer_with_param(trainable_weights,
-                                                                  trainable_bias,
-                                                                  trainable_threshold)
-        hessian_cfg = self.gptq_config.hessian_weights_config
-
-        self.use_sample_layer_attention = hessian_cfg and hessian_cfg.per_sample
-        if self.use_sample_layer_attention:
-            # normalization is currently not supported, make sure the config reflects it.
-            if hessian_cfg.norm_scores or hessian_cfg.log_norm or hessian_cfg.scale_log_norm:
-                raise NotImplementedError()
-            self.train_dataloader = self._prepare_train_dataloader_sla(representative_data_gen)
-        else:
-            self.train_dataloader = self._prepare_train_dataloader_for_non_sla(representative_data_gen)
-
-        self.reg_func = get_regularization(self.gptq_config, _get_total_grad_steps, PytorchSoftQuantizerRegularization, PytorchLinearAnnealingScheduler)
 
     def _prepare_train_dataloader_sla(self, data_gen_fn: Callable[[], Generator]) -> DataLoader:
         """