def trace_gptq_target_blocks(
    module, module_types=(torch.nn.ModuleList, torch.nn.Sequential)
):
    """Search transformer stacked structures, which is critical in LLMs and GPTQ execution.

    Args:
        module: torch.nn.Module
        module_types: List of torch.nn.Module.

    Returns:
        gptq_related_blocks = {
            "embeddings": {}, # Dict embedding layers before transformer stack module
            "transformers_pre": {}, # TODO
            "transformers_name": "", # Str, LLMs' transformer stack module name
            "transformers": # torch.nn.ModuleList. LLMs' transformer stack module
            "transformers": {}, # TODO
        }
    """
    if type(module).__name__ == "MixFormerSequentialForCausalLM":
        gptq_related_blocks = {
            "embeddings": {},
            "transformers_pre": {},
            "transformers_name": "",
            "transformers": [],
            "transformers_post": {},
        }
        for n, m in module.named_modules():
            if type(m) in module_types:
                gptq_related_blocks["transformers_name"] = n
                gptq_related_blocks["transformers"] = m
                break
            else:
                continue
        for n, m in gptq_related_blocks["transformers"][0].named_modules():
            if is_leaf(m):
                gptq_related_blocks["embeddings"][n] = m
        gptq_related_blocks["transformers"] = gptq_related_blocks["transformers"][1:-1]
    else:
        gptq_related_blocks = {
            "embeddings": {},
            "transformers_pre": {},
            "transformers_name": "",
            "transformers": [],
            "transformers_post": {},
        }
        for n, m in module.named_modules():
            if type(m) in module_types:
                gptq_related_blocks["transformers_name"] = n
                gptq_related_blocks["transformers"] = m
                return gptq_related_blocks
            else:
                if is_leaf(m):
                    gptq_related_blocks["embeddings"][n] = m
    return gptq_related_blocks
Search transformer stacked structures, which is critical in LLMs and GPTQ execution.

    Args:
        module: torch.nn.Module
        module_types: List of torch.nn.Module.

    Returns:
        gptq_related_blocks = {
            "embeddings": {}, # Dict embedding layers before transformer stack module
            "transformers_pre": {}, # TODO
            "transformers_name": "", # Str, LLMs' transformer stack module name
            "transformers": # torch.nn.ModuleList. LLMs' transformer stack module
            "transformers": {}, # TODO
        }
    
intel_extension_for_pytorch/quantization/_GPTQ/gptq/model_utils.py