code: def load_args_from_checkpoint(args):
      # Read Llama args.
      llama_args_path = os.path.join(args.load, "config.json")
      with open(llama_args_path) as f:
          llama_args = json.load(f)
  
      # Update Megatron args.
      args.seq_length = 4096
      args.max_position_embeddings = 4096
      args.hidden_size = llama_args["hidden_size"]
      args.num_attention_heads = llama_args["num_attention_heads"]
      args.num_layers = llama_args["num_hidden_layers"]
      args.global_batch_size = 1024
      args.norm_epsilon = llama_args["rms_norm_eps"]
      args.iteration = 1  # '0', 'release' don't work
      args.add_position_embedding = False
      args.use_rotary_position_embeddings = True
      args.swiglu = True
      args.tokenizer_type = "Llama2Tokenizer"
      args.fp16 = True
      args.normalization = "RMSNorm"
      args.add_bias_linear = False
      args.untie_embeddings_and_output_weights = True
      args.vocab_size = llama_args["vocab_size"]
      args.padded_vocab_size = llama_args["vocab_size"]
      args.llama = llama_args
      args.ffn_hidden_size = llama_args["intermediate_size"]
  
      if "num_key_value_heads" in llama_args:
          args.group_query_attention = True
          args.num_query_groups = llama_args["num_key_value_heads"]
  
file_path: megatron-lm/tools/checkpoint/loader_llama2_hf.py
imports: import json
  import os
  import sys
  import torch
  import transformers
  from tqdm import tqdm
  import types
label_desc: Use image node_id 3 for calling a global function with example usage: load_args_from_checkpoint(args) without return types