docstr: Custom initialization of Megatron.
docstr_tok: ["Custom","initialization","of","Megatron","."]
code: def init_megatron(cls, workdir):
      """Custom initialization of Megatron."""
  
      # Load args.
      args_path = get_args_path(workdir)
      assert os.path.exists(
          args_path
      ), "args.json not found in workdir."
      with open(args_path) as f:
          cls.args = types.SimpleNamespace(**json.load(f))
          cls.args.retro_workdir = workdir  # just in case workdir moved
          cls.args.rank = 0  # override env
          cls.args.world_size = 1  # override env
          cls.args.params_dtype = cls.parse_dtype_str(
              cls.args.params_dtype
          )
          cls.args.retro_verify_neighbor_count = False
  
      set_global_variables(cls.args)
      set_retro_args(cls.args)
      _initialize_distributed()
      _set_random_seed(
          cls.args.seed, cls.args.data_parallel_random_init
      )
      _compile_dependencies()
  
imports: import json
  import numpy
  import os
  import torch
  import types
  from megatron.global_vars import set_global_variables, set_retro_args
  from megatron.initialize import initialize_megatron, _initialize_distributed, _set_random_seed, _compile_dependencies
  from tools.retro.db.utils import get_indexed_dataset_infos, get_merged_train_dataset
  from tools.retro.main import add_retro_args
  from tools.retro.query.retro_dataset import get_retro_datasets
  from tools.retro.utils import get_args_path, get_bert_tokenizer, get_gpt_tokenizer
label_desc: Use image node_id 2 for calling the retro obj's underlying member method code with example usage: obj.init_megatron(cls, workdir) without return types