code: def quantize(self, inputs, qkv=True, count=1, parallel_dim=0):
      if not self.q_int8 or not qkv:
          inputs = torch.nn.Parameter(inputs, requires_grad=False)
          inputs.scale = torch.empty(1)
          return inputs
      q_range = 2**self.num_bits
      num_groups = (
          self.num_groups
          if self.num_groups > 0
          else inputs.shape[0] // self.group_size
      )
      inputs = inputs.to(get_accelerator().current_device_name())
      input_flat = inputs.reshape(num_groups, -1).contiguous()
      input_min = torch.min(input_flat, dim=1, keepdim=True)[0].float()
      input_max = torch.max(input_flat, dim=1, keepdim=True)[0].float()
      scale = (
          torch.max(input_min.abs(), input_max.abs()) * 2.0 / (q_range)
      )
      input_flat = (
          (input_flat / scale)
          .round()
          .clamp(-q_range // 2, q_range // 2 - 1)
      )
      inputs_q = (
          input_flat.reshape(inputs.shape).to(torch.int8).contiguous()
      )
      out = torch.nn.Parameter(inputs_q, requires_grad=False)
      inputs_split = inputs.split(
          inputs.shape[parallel_dim] // 2, dim=parallel_dim
      )
      input_flat = [
          inputs_split[i].reshape(num_groups, -1).contiguous()
          for i in range(2)
      ]
      input_min = [
          torch.min(input_flat[i], dim=1, keepdim=True)[0].float()
          for i in range(2)
      ]
      input_max = [
          torch.max(input_flat[i], dim=1, keepdim=True)[0].float()
          for i in range(2)
      ]
      scale1 = [
          (
              torch.max(input_min[i].abs(), input_max[i].abs())
              * 2.0
              / (q_range)
          )
          .squeeze()
          .unsqueeze(0)
          for i in range(2)
      ]
  
      out.scale = (
          torch.cat(
              [scale.squeeze().unsqueeze(0), scale1[0], scale1[1]],
              dim=0,
          )
          .reshape(num_groups, -1)
          .contiguous()
      )
      return out
  
imports: import os
  import torch
  import tqdm
  import deepspeed
  import deepspeed.ops.transformer
  from deepspeed.ops.transformer.inference.diffusers_attention import DeepSpeedDiffusersAttention
  from deepspeed.ops.transformer.inference.diffusers_transformer_block import DeepSpeedDiffusersTransformerBlock
  from deepspeed.ops.transformer.inference.diffusers_2d_transformer import Diffusers2DTransformerConfig
  from deepspeed.accelerator import get_accelerator
  from .replace_policy import replace_policies, generic_policies
  from .auto_tp import AutoTP, ReplaceWithTensorSlicing, Loading
  from deepspeed import comm
  from deepspeed.module_inject.tp_shard import set_num_kv_heads
  from .load_checkpoint import load_model_with_checkpoint
  import time
  from .utils import policy_to_ds_container
  import gc
  from ..pipe import PipelineModule
  import re
label_desc: Use image node_id 2 for calling the GroupQuantizer obj's underlying member method code with example usage: obj.quantize(inputs, qkv, count, parallel_dim) and returns: out, inputs