code: def fix_query_key_value_ordering(
      param, checkpoint_version, num_splits, num_heads, hidden_size
  ):
      # Permutes layout of param tensor to [num_splits * num_heads * hidden_size, :]
      # for compatibility with later versions of NVIDIA Megatron-LM.
      # The inverse operation is performed inside Megatron-LM to read checkpoints:
      # https://github.com/NVIDIA/Megatron-LM/blob/v2.4/megatron/checkpointing.py#L209
      # If param is the weight tensor of the self-attention block, the returned tensor
      # will have to be transposed one more time to be read by HuggingFace BERT.
      input_shape = param.size()
      if checkpoint_version == 1.0:
          # version 1.0 stores [num_heads * hidden_size * num_splits, :]
          saved_shape = (
              num_heads,
              hidden_size,
              num_splits,
          ) + input_shape[1:]
          param = param.view(*saved_shape)
          param = param.transpose(0, 2)
          param = param.transpose(1, 2).contiguous()
      elif checkpoint_version >= 2.0:
          # other versions store [num_heads * num_splits * hidden_size, :]
          saved_shape = (
              num_heads,
              num_splits,
              hidden_size,
          ) + input_shape[1:]
          param = param.view(*saved_shape)
          param = param.transpose(0, 1).contiguous()
      param = param.view(*input_shape)
      return param
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/src/transformers/models/megatron_bert/convert_megatron_bert_checkpoint.py
imports: import argparse
  import os
  import re
  import zipfile
  import torch
  from transformers import MegatronBertConfig
label_desc: Use image node_id 2 for calling a global function with example usage: fix_query_key_value_ordering(param, checkpoint_version, num_splits, num_heads, hidden_size) and returns: param