def pre_quantization(self):
        """Prepare input calibration data and other attributes which are critical for gptq execution."""

        # critical: hooker function which collects inputs
        def forward(layer, *args, **kwargs):
            self.cache_key_arguments["i"] += 1
            for arg in kwargs:
                # each outputs can be different shape, hence also use list to store
                if isinstance(kwargs[arg], torch.Tensor) or arg == "alibi":
                    if self.cache_key_arguments.get(arg, None) is None:
                        self.cache_key_arguments[arg] = []
                    self.cache_key_arguments[arg].append(kwargs[arg])
                continue
            # copy positional arguments, positional arguments are sensitive for their order, be cautious!
            # Most models in HF has avoid this, but some models still use positional arguments other than
            # hidden_states, chatglm2-6b etc.
            for idx, item in enumerate(args):
                if (idx + 1) > len(self.cache_positional_arguments):
                    # initialize
                    self.cache_positional_arguments.append([])
                self.cache_positional_arguments[idx].append(item)
            raise ValueError

        # Step1: fetch the embeddings and other layers before the transformer stack.
        for embedding_name, embedding_layer in self.gptq_related_blocks[
            "embeddings"
        ].items():
            embedding_layer = embedding_layer.to(self.device)

        # Step2: modify the first transformer block's forward function to obtain inputs for calibration
        self.gptq_related_blocks["transformers"][0] = self.gptq_related_blocks[
            "transformers"
        ][0].to(self.device)
        forward_cache = self.gptq_related_blocks["transformers"][0].forward
        self.gptq_related_blocks["transformers"][0].forward = partial(
            forward, self.gptq_related_blocks["transformers"][0]
        )

        # Step3: run forward to obtain calibration datasets
        logger.info("Collecting calibration inputs...")
        for batch in tqdm(self.dataloader):
            batch = move_input_to_device(batch, self.device)
            try:
                if isinstance(batch, tuple) or isinstance(batch, list):
                    self.model(batch[0])
                elif isinstance(batch, dict):
                    self.model(**batch)
                else:
                    self.model(batch)
            except ValueError:
                pass
        # output inp data shape
        logger.info("All calibration data's shape =>")
        # check all hidden_states shape
        try:
            for hidden_states in self.cache_positional_arguments[0]:
                logger.info(hidden_states.shape)
        except Exception:
            pass
        logger.info("Done.")

        # Step 4: restore original forward function, relocate layers back to cpu.
        self.gptq_related_blocks["transformers"][0].forward = forward_cache
        self.gptq_related_blocks["transformers"][0] = self.gptq_related_blocks[
            "transformers"
        ][0].cpu()
        for embedding_name, embedding_layer in self.gptq_related_blocks[
            "embeddings"
        ].items():
            embedding_layer.to(self.device)
        torch.cuda.empty_cache()
        # end
        logger.info("GPTQ quantization prepared.")
Prepare input calibration data and other attributes which are critical for gptq execution.
intel_extension_for_pytorch/quantization/_GPTQ/gptq/gptq.py