code: def convert_megatron_checkpoint(sd_megatron, config):
      """
      Converts a Megatron checkpoint to a HuggingFace GPT-SW3 checkpoint.
      """
      n_positions = config.n_positions
      layers = config.n_layer
      vocab_size = config.vocab_size
      heads = config.n_head
      hidden_size_per_head = config.n_embd // config.n_head
  
      word_embeddings = sd_megatron[
          "model.language_model.embedding.word_embeddings.weight"
      ][:vocab_size, :]
      sd_hf = {
          "transformer.wte.weight": word_embeddings,
          "transformer.wpe.weight": sd_megatron[
              "model.language_model.embedding.position_embeddings.weight"
          ],
          "transformer.ln_f.weight": sd_megatron[
              "model.language_model.encoder.final_layernorm.weight"
          ],
          "transformer.ln_f.bias": sd_megatron[
              "model.language_model.encoder.final_layernorm.bias"
          ],
      }
  
      pf = "model.language_model.encoder.layers."
      for i in range(layers):
          causal_mask = torch.tril(
              torch.ones((n_positions, n_positions), dtype=torch.bool)
          )
          causal_mask = causal_mask.view(1, 1, n_positions, n_positions)
          sd_hf[f"transformer.h.{i}.attn.bias"] = causal_mask
          sd_hf[f"transformer.h.{i}.attn.masked_bias"] = torch.tensor(
              -1e4, dtype=torch.bfloat16
          )
  
          sd_hf[f"transformer.h.{i}.ln_1.weight"] = sd_megatron[
              f"{pf}{i}.input_layernorm.weight"
          ]
          sd_hf[f"transformer.h.{i}.ln_1.bias"] = sd_megatron[
              f"{pf}{i}.input_layernorm.bias"
          ]
  
          val1 = sd_megatron[
              f"{pf}{i}.self_attention.query_key_value.weight"
          ]
          val1 = fix_query_key_value_ordering(
              val1, 3, heads, hidden_size_per_head
          )
          sd_hf[
              f"transformer.h.{i}.attn.c_attn.weight"
          ] = val1.transpose(0, 1).contiguous()
  
          val2 = sd_megatron[
              f"{pf}{i}.self_attention.query_key_value.bias"
          ]
          val2 = fix_query_key_value_ordering(
              val2, 3, heads, hidden_size_per_head
          )
          sd_hf[f"transformer.h.{i}.attn.c_attn.bias"] = val2
  
          sd_hf[f"transformer.h.{i}.attn.c_proj.weight"] = sd_megatron[
              f"{pf}{i}.self_attention.dense.weight"
          ].transpose(0, 1)
          sd_hf[f"transformer.h.{i}.attn.c_proj.bias"] = sd_megatron[
              f"{pf}{i}.self_attention.dense.bias"
          ]
          sd_hf[f"transformer.h.{i}.ln_2.weight"] = sd_megatron[
              f"{pf}{i}.post_attention_layernorm.weight"
          ]
          sd_hf[f"transformer.h.{i}.ln_2.bias"] = sd_megatron[
              f"{pf}{i}.post_attention_layernorm.bias"
          ]
          sd_hf[f"transformer.h.{i}.mlp.c_fc.weight"] = sd_megatron[
              f"{pf}{i}.mlp.dense_h_to_4h.weight"
          ].transpose(0, 1)
          sd_hf[f"transformer.h.{i}.mlp.c_fc.bias"] = sd_megatron[
              f"{pf}{i}.mlp.dense_h_to_4h.bias"
          ]
          sd_hf[f"transformer.h.{i}.mlp.c_proj.weight"] = sd_megatron[
              f"{pf}{i}.mlp.dense_4h_to_h.weight"
          ].transpose(0, 1)
          sd_hf[f"transformer.h.{i}.mlp.c_proj.bias"] = sd_megatron[
              f"{pf}{i}.mlp.dense_4h_to_h.bias"
          ]
  
      # For LM head, transformers' wants the matrix to weight embeddings.
      sd_hf["lm_head.weight"] = word_embeddings
  
      return sd_hf
  
file_path: transformers/src/transformers/models/gpt_sw3/convert_megatron_to_pytorch.py
imports: import argparse
  import os
  from os.path import isfile
  import torch
  from transformers import GPT2Config
label_desc: Use image node_id 3 for calling a global function with example usage: convert_megatron_checkpoint(sd_megatron, config) and returns: sd_hf