code: def load_col(config, prefix, weights, bias):
      assert bias == False, NotImplementedError
      assert config.quantize != "gptq", NotImplementedError
      slice_ = weights._get_slice(f"{prefix}.weight")
      rank = weights.process_group.rank()
      size = weights.process_group.size()
  
      h3, h = slice_.get_shape()
      block_size = h // size
  
      q_part = slice_[rank * block_size : (rank + 1) * block_size]
      k_part = slice_[
          h + rank * block_size : h + (rank + 1) * block_size
      ]
      v_part = slice_[
          2 * h + rank * block_size : 2 * h + (rank + 1) * block_size
      ]
  
      weight = torch.cat([q_part, k_part, v_part], dim=0)
      if weight.dtype != torch.int32:
          weight = weight.to(dtype=weights.dtype)
      weight = weight.to(device=weights.device)
      bias = None
      linear = get_linear(weight, bias, config.quantize)
      return TensorParallelColumnLinear(linear)
  
imports: import math
  import os
  import warnings
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn
  import torch.nn.functional
  from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
  from einops import rearrange
  from packaging import version
  from text_generation_server.utils.layers import TensorParallelEmbedding, TensorParallelColumnLinear, TensorParallelRowLinear, TensorParallelHead, get_linear
label_desc: Use image node_id 1 for calling a global function with example usage: load_col(config, prefix, weights, bias) and returns: TensorParallelColumnLinear