docstr: Up to now we found 3 Q/K/V parameter formats in different Megatron checkpoint versions:
  
  1. version 0, there is no version information saved in checkpoint.
      format: [(3 * np * hn), h]
  2. version 1.0
      format: [(np * hn * 3), h]
  3. version 2.0
      format: [(np * 3 * hn), h]
  
  h: hidden size
  n: number of attention heads
  p: number of model parallel partitions
  np: n/p
  hn: h/n
docstr_tok: ["Up","to","now","we","found","3","Q\/K\/V","parameter","formats","in","different","Megatron","checkpoint","versions",":","1.","version","0",",","there","is","no","version","information","saved","in","checkpoint",".","format",":","[","(","3","*","np","*","hn",")",",","h","]","2.","version","1.0","format",":","[","(","np","*","hn","*","3",")",",","h","]","3.","version","2.0","format",":","[","(","np","*","3","*","hn",")",",","h","]","h",":","hidden","size","n",":","number","of","attention","heads","p",":","number","of","model","parallel","partitions","np",":","n\/p","hn",":","h\/n"]
code: def split_query_key_value(
      self, param, num_to_split, offset, ckpt_ver
  ):
      """
      Up to now we found 3 Q/K/V parameter formats in different Megatron checkpoint versions:
  
      1. version 0, there is no version information saved in checkpoint.
          format: [(3 * np * hn), h]
      2. version 1.0
          format: [(np * hn * 3), h]
      3. version 2.0
          format: [(np * 3 * hn), h]
  
      h: hidden size
      n: number of attention heads
      p: number of model parallel partitions
      np: n/p
      hn: h/n
      """
  
      new_qkv = None
      if ckpt_ver == 0:
          # [(3 * np * hn), h]
          assert param.shape[0] % 3 == 0
          size_qkv = param.shape[0] // 3
          split_tensors = torch.split(param, size_qkv, dim=0)
  
          assert split_tensors[0].shape[0] % num_to_split == 0
          split_size = split_tensors[0].shape[0] // num_to_split
  
          tensors = []
          for i in range(3):
              tensors.append(
                  torch.split(split_tensors[i], split_size, dim=0)[
                      offset
                  ]
              )
          new_qkv = torch.cat(tensors, axis=0)
      elif ckpt_ver == 1.0 or ckpt_ver == 2.0:
          # [(np * hn * 3), h] or [(np * 3 * hn), h]
          assert param.shape[0] % num_to_split == 0
          size_qkv = param.shape[0] // num_to_split
          split_tensors = torch.split(param, size_qkv, dim=0)
          new_qkv = split_tensors[offset]
      else:
          assert (
              False
          ), f"checkpoint version: {ckpt_ver} is not supported"
  
      return new_qkv
  
imports: import torch
  import os
  import copy
  import collections
  import json
  from abc import ABC, abstractmethod
  from deepspeed.utils import logger
  from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine
  from .weight_quantizer import WeightQuantization
label_desc: Use image node_id 3 for calling the MegatronSDLoader obj's underlying member method code with example usage: obj.split_query_key_value(param, num_to_split, offset, ckpt_ver) and returns: new_qkv