code: def test_gptq_linear():
      infeature = 1024
      outfeature = 1024
      group_size = 128
      wbits = 4
  
      inps = (
          torch.ones(1, 1, infeature)
          .to(torch.float16)
          .to(torch.cuda.current_device())
      )
      batch_inps = (
          torch.randn(1, 16, infeature)
          .to(torch.float16)
          .to(torch.cuda.current_device())
      )
  
      device = torch.device("cuda:0")
  
      linear_class = dynamically_import_QuantLinear(
          use_triton=False,
          desc_act=False,
          group_size=group_size,
          bits=wbits,
      )
  
      linear = linear_class(
          bits=4,
          group_size=group_size,
          infeatures=infeature,
          outfeatures=outfeature,
          bias=False,
      )
  
      torch.manual_seed(42)
  
      linear.qweight = torch.randint(
          -100, 100, size=linear.qweight.shape, dtype=torch.int32
      )
      linear.scales = linear.scales + 0.002
  
      linear = linear.to(device)
  
      cai_linear = CaiQuantLinear(
          wbits, group_size, infeature, outfeature, True
      )
      cai_linear.qweight.data.copy_(linear.qweight)
      cai_linear.scales = cai_linear.scales + 0.002
      cai_linear = cai_linear.to(device)
  
      linear = autogptq_post_init(linear, use_act_order=False)
  
      max_inner_outer_dim = max(infeature, outfeature)
      max_dq_buffer_size = linear.infeatures * linear.outfeatures
      max_input_len = 2048
      buffers = {
          "temp_state": torch.zeros(
              (max_input_len, max_inner_outer_dim),
              dtype=torch.float16,
              device=device,
          ),
          "temp_dq": torch.zeros(
              (1, max_dq_buffer_size),
              dtype=torch.float16,
              device=device,
          ),
      }
  
      prepare_buffers(device, buffers["temp_state"], buffers["temp_dq"])
  
      # Using the default from exllama repo here.
      matmul_recons_thd = 8
      matmul_fused_remap = False
      matmul_no_half2 = False
      set_tuning_params(
          matmul_recons_thd, matmul_fused_remap, matmul_no_half2
      )
  
      with torch.no_grad():
          gptq_out = linear(inps)
          batch_gptq_out = linear(batch_inps)
          torch.cuda.synchronize()
          cai_out = cai_linear(inps)
          torch.cuda.synchronize()
  
          batch_cai_out = cai_linear(batch_inps)
          torch.cuda.synchronize()
  
      assert torch.allclose(cai_out, gptq_out, rtol=1e-01, atol=1e-01)
      assert torch.allclose(
          batch_cai_out, batch_gptq_out, rtol=1e-01, atol=1e-01
      )
  
file_path: colossalai/tests/test_gptq/test_gptq_linear.py
imports: import math
  import time
  import numpy
  import pytest
  import torch
  import torch.nn
  import transformers
  from packaging import version
  import warnings
label_desc: Use image node_id 2 for calling a global function with example usage: test_gptq_linear() without return types