# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.

import json
import os
import sys
import types

import torch

def add_arguments(parser):
    group = parser.add_argument_group(title='Megatron loader')

    group.add_argument('--true-vocab-size', type=int, default=None,
                       help='original size of vocab, if specified will trim padding from embedding table.')
    group.add_argument('--vocab-file', type=str, default=None,
                       help='Path to the vocab file. If specified will use this to get vocab size and '
                       'trim padding from the embedding table.')
    group.add_argument('--megatron-path', type=str, default=None,
                       help='Base directory of deepspeed repository')

def _load_checkpoint(queue, args):

    # Search in directory above this
    sys.path.append(os.path.abspath(
        os.path.join(os.path.dirname(__file__),
                     os.path.pardir)))
    if args.megatron_path is not None:
        sys.path.insert(0, args.megatron_path)

    try:
        from megatron.arguments import parse_args, validate_args
        from megatron.global_vars import set_args, set_global_variables
        from megatron.checkpointing import load_args_from_checkpoint, load_checkpoint
        from megatron.model import module
        from megatron.core import mpu
        from megatron.core.enums import ModelType
        from megatron import fused_kernels
    except ModuleNotFoundError:
        print("Unable to import Megatron, please specify the path to Megatron using --megatron-path. Exiting.")


# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.

import json
import os
import sys
import types

import torch

def add_arguments(parser):
    group = parser.add_argument_group(title='Megatron loader')

    group.add_argument('--true-vocab-size', type=int, default=None,
                       help='original size of vocab, if specified will trim padding from embedding table.')
    group.add_argument('--vocab-file', type=str, default=None,
                       help='Path to the vocab file. If specified will use this to get vocab size and '
                       'trim padding from the embedding table.')
    group.add_argument('--megatron-path', type=str, default=None,
                       help='Base directory of deepspeed repository')

def _load_checkpoint(queue, args):

    # Search in directory above this
    sys.path.append(os.path.abspath(
        os.path.join(os.path.dirname(__file__),
                     os.path.pardir)))
    if args.megatron_path is not None:
        sys.path.insert(0, args.megatron_path)

    try:
        from megatron.arguments import parse_args, validate_args
        from megatron.global_vars import set_args, set_global_variables
        from megatron.checkpointing import load_args_from_checkpoint, load_checkpoint
        from megatron.model import module
        from megatron.core import mpu
        from megatron.core.enums import ModelType
        from megatron import fused_kernels
    except ModuleNotFoundError:
        print("Unable to import Megatron, please specify the path to Megatron using --megatron-path. Exiting.")
        queue.put("exit")
        exit(1)

    # We want all arguments to come from us
    sys.argv = ['script.py',
                '--no-masked-softmax-fusion',
                '--no-bias-gelu-fusion',
                '--no-bias-dropout-fusion',
                '--no-async-tensor-model-parallel-allreduce',
                '--use-cpu-initialization',
                '--micro-batch-size', '1',
                '--no-load-optim',
                '--no-load-rng',
                '--no-save-optim',
                '--no-save-rng',
                '--no-initialization',
                '--load', args.load_dir,
                '--load-iteration', str(args.load_iteration)
                ]

    margs = parse_args()
    margs, checkpoint_args = load_args_from_checkpoint(margs)

    # Arguments do sanity checks on the world size, but we don't care,
    # so trick it into thinking we are plenty of processes
    margs.world_size = margs.tensor_model_parallel_size * margs.pipeline_model_parallel_size

    margs = validate_args(margs)

    def check_for_arg(arg_name, default=None):
        if getattr(margs, arg_name, None) is None:
            if default is not None:
                setattr(margs, arg_name, default)
            else:
                print(f"Checkpoint does not specify the argument {arg_name}. Exiting.")
                print(f"Arguments: {margs}")
                queue.put("exit")
                exit(1)

    check_for_arg('tensor_model_parallel_size')
    check_for_arg('pipeline_model_parallel_size')
    check_for_arg('num_layers')
    check_for_arg('hidden_size')
    check_for_arg('seq_length')
    check_for_arg('num_attention_heads')
    check_for_arg('max_position_embeddings')
    check_for_arg('position_embedding_type')
    check_for_arg('tokenizer_type')
    check_for_arg('iteration')
    check_for_arg('bert_binary_head')
    check_for_arg('disable_bias_linear', False)
    check_for_arg('params_dtype')
    check_for_arg('swiglu', False)

    # Determine how to make our models
    if args.model_type == 'GPT':
        from pretrain_gpt import model_provider
        margs.model_type = ModelType.encoder_or_decoder
    elif args.model_type == 'BERT':
        from pretrain_bert import model_provider
        margs.model_type = ModelType.encoder_or_decoder
    elif args.model_type == 'EarlyExitGPT':
        from pretrain_early_exit_gpt import model_provider
        margs.model_type = ModelType.encoder_or_decoder
    else:
        raise Exception(f'unrecognized model type: {args.model_type}')

    # supress warning about torch.distributed not being initialized
    module.MegatronModule.embedding_warning_printed = True

    consumed_train_samples = None
    consumed_valid_samples = None
    def get_models(count, dtype):
        nonlocal consumed_train_samples
        nonlocal consumed_valid_samples
        model_array_len = margs.virtual_pipeline_model_parallel_size
        if model_array_len is None:
            model_array_len = 1
        models = [[] for _ in range(model_array_len)]
        pre_process = mpu.is_pipeline_first_stage()
        post_process = mpu.is_pipeline_last_stage()
        for rank in range(count):
            mpu.set_tensor_model_parallel_rank(rank)
            if margs.virtual_pipeline_model_parallel_size is not None:
                model_ = []
                for i in range(margs.virtual_pipeline_model_parallel_size):
                    mpu.set_virtual_pipeline_model_parallel_rank(i)
                    # Set pre_process and post_process only after virtual rank is set.
                    pre_process = mpu.is_pipeline_first_stage()
                    post_process = mpu.is_pipeline_last_stage()
                    this_model = model_provider(
                        pre_process=pre_process,
                        post_process=post_process
                    ).to(dtype)
                    model_.append(this_model)
            else:
                pre_process = mpu.is_pipeline_first_stage()
                post_process = mpu.is_pipeline_last_stage()
                model_rank = 0
                model_ = [model_provider(pre_process, post_process).to(dtype)]
            margs.consumed_train_samples = 0
            margs.consumed_valid_samples = 0
            load_checkpoint(model_, None, None)

            if consumed_train_samples is not None:
                assert(margs.consumed_train_samples == consumed_train_samples)
            else:
                consumed_train_samples = margs.consumed_train_samples
            if consumed_valid_samples is not None:
                assert(margs.consumed_valid_samples == consumed_valid_samples)
            else:
                consumed_valid_samples = margs.consumed_valid_samples
            for vp_rank in range(model_array_len):
                models[vp_rank].append(model_[vp_rank])
        return models

    set_global_variables(margs, build_tokenizer=False)
    mpu.set_tensor_model_parallel_world_size(margs.tensor_model_parallel_size)
    mpu.set_pipeline_model_parallel_world_size(margs.pipeline_model_parallel_size)
    mpu.set_virtual_pipeline_model_parallel_world_size(margs.virtual_pipeline_model_parallel_size)
    fused_kernels.load(margs)

    # Get true (non-padded) vocab size
    if args.true_vocab_size is not None:
        true_vocab_size = args.true_vocab_size
    elif args.vocab_file is not None:
        vocab = json.load(open(args.vocab_file))
        true_vocab_size = len(vocab)
        if args.true_vocab_size is not None and true_vocab_size != args.true_vocab_size:
            print("Both --true-vocab-size and --vocab-file specified and the vocab size does not match, aborting.")
            queue.put("exit")
            exit(1)
    else:
        true_vocab_size = None

    # short aliases
    tp_size = margs.tensor_model_parallel_size
    pp_size = margs.pipeline_model_parallel_size
    vp_size = margs.virtual_pipeline_model_parallel_size
    if vp_size is None:
        vp_size = 1

    # Layernorm has bias; RMSNorm does not.
    if hasattr(checkpoint_args, 'normalization'):
        norm_has_bias = checkpoint_args.normalization == "LayerNorm"
    else:
        # older models only supported LayerNorm
        norm_has_bias = True

    # metadata
    md = types.SimpleNamespace()
    md.model_type = args.model_type
    md.num_layers = margs.num_layers
    md.hidden_size = margs.hidden_size
    md.seq_length = margs.seq_length
    md.num_attention_heads = margs.num_attention_heads
    md.max_position_embeddings = margs.max_position_embeddings
    md.tokenizer_type = margs.tokenizer_type
    md.iteration = margs.iteration
    md.params_dtype = margs.params_dtype
    md.bert_binary_head = margs.bert_binary_head
    md.output_layer = margs.untie_embeddings_and_output_weights
    md.position_embedding_type = margs.position_embedding_type
    md.linear_bias = margs.add_bias_linear
    md.norm_has_bias = norm_has_bias
    md.swiglu = margs.swiglu
    md.previous_tensor_parallel_size = margs.tensor_model_parallel_size
    md.previous_pipeline_parallel_size = margs.pipeline_model_parallel_size
    md.true_vocab_size = true_vocab_size
    md.make_vocab_size_divisible_by = margs.make_vocab_size_divisible_by
    md.exit_layer_nums = margs.exit_layer_nums if hasattr(margs, 'exit_layer_nums') else []
    md.exit_layer_weight = margs.exit_layer_weight if hasattr(margs, 'exit_layer_weight') else []
    md.use_exit_mlp = margs.use_exit_mlp if hasattr(margs, 'use_exit_mlp') else False
    md.use_exit_block = margs.use_exit_block if hasattr(margs, 'use_exit_block') else False
    md.use_exit_norm = margs.use_exit_norm if hasattr(margs, 'use_exit_norm') else False
    md.untie_exit_output_weights = margs.untie_exit_output_weights if hasattr(margs, 'untie_exit_output_weights') else False
    md.pre_exit = margs.pre_exit
    md.checkpoint_args = checkpoint_args

    # Get first pipe stage
    mpu.set_pipeline_model_parallel_rank(0)
    if len(md.exit_layer_nums) > 0:
        layer_per_stage = md.num_layers / margs.pipeline_model_parallel_size
        mpu.set_early_exit_layer_nums(list(filter(lambda x: 0 < x <= layer_per_stage, md.exit_layer_nums)))
        mpu.set_early_exit_stages(list(set(map(lambda layer_num:  int((layer_num - 1) // layer_per_stage), md.exit_layer_nums))))
    all_models = [get_models(tp_size, md.params_dtype)]
    models = all_models[0][0]

    md.consumed_train_samples = consumed_train_samples
    md.consumed_valid_samples = consumed_valid_samples
    queue.put(md)

    def queue_put(name, msg):
        print(f"sending {name}")
        msg["name"] = name
        queue.put(msg)

    # Send embeddings
    message = {
        "word embeddings": torch.cat(
            [models[tp_rank].language_model.embedding.word_embeddings.weight.data for tp_rank in range(tp_size)],
            dim = 0)
    }
    if md.position_embedding_type == 'learned_absolute':
        message["position embeddings"] = models[0].language_model.embedding.position_embeddings.weight.data
    else:
        assert not hasattr(models[0].language_model.embedding, 'position_embeddings')

    queue_put("embeddings", message)

    total_layer_num = 0
    for vp_rank in range(vp_size):
        mpu.set_virtual_pipeline_model_parallel_rank(vp_rank)
        for pp_rank in range(pp_size):
            if pp_rank > 0:
                mpu.set_pipeline_model_parallel_rank(pp_rank)
                if len(md.exit_layer_nums) > 0:
                    mpu.set_early_exit_layer_nums(list(filter(lambda x: (layer_per_stage * pp_rank) < x <= (layer_per_stage * (pp_rank + 1)), md.exit_layer_nums)))
                if vp_rank == 0:
                    all_models.append(get_models(tp_size, md.params_dtype))
            models = all_models[pp_rank][vp_rank]
            for layer_id in range(len(models[0].language_model.encoder.layers)):
                message = {}
                # Get non-parallel tensors from tp_rank 0
                layer = models[0].language_model.encoder.layers[layer_id]
                layer_num = layer.layer_number
                has_early_exit = layer_num in md.exit_layer_nums
                use_exit_mlp = has_early_exit and md.use_exit_mlp
                use_exit_block = has_early_exit and md.use_exit_block
                use_exit_norm = has_early_exit and md.use_exit_norm
                message["input norm weight"] = layer.input_norm.weight.data
                if norm_has_bias:
                    message["input norm bias"] = layer.input_norm.bias.data
                message["post norm weight"] = layer.post_attention_norm.weight.data
                if norm_has_bias:
                    message["post norm bias"] = layer.post_attention_norm.bias.data
                if use_exit_norm:
                    message["exit norm weight"] = layer.exit_norm.weight.data
                    if norm_has_bias:
                        message['exit norm bias'] = layer.exit_norm.bias.data
                if md.linear_bias:
                    message["dense bias"] = layer.self_attention.dense.bias.data
                    if use_exit_mlp:
                        message["mlp l1 bias"] = layer.mlp.trunk.dense_4h_to_h.bias.data
                        message["mlp l1 exit bias"] = layer.mlp.branch.dense_4h_to_h.bias.data
                    else:
                        message["mlp l1 bias"] = layer.mlp.dense_4h_to_h.bias.data
                if use_exit_block:
                    message["exit block input norm weight"] = layer.exit_block.input_norm.weight.data
                    message["exit block post norm weight"] = layer.exit_block.post_attention_norm.weight.data
                    if norm_has_bias:
                        message["exit block input norm bias"] = layer.exit_block.input_norm.bias.data
                        message["exit block post norm bias"] = layer.exit_block.post_attention_norm.bias.data
                    if md.linear_bias:
                        message["exit block dense bias"] = layer.exit_block.self_attention.dense.bias.data
                        message["exit block mlp l1 bias"] = layer.exit_block.mlp.dense_4h_to_h.bias.data

                # Grab all parallel tensors for this layer
                qkv_weight = []
                qkv_bias = []
                dense_weight = []
                mlp_l0_weight = []
                mlp_l0_bias = []
                mlp_l1_weight = []
                mlp_l0_exit_weight = []
                mlp_l0_exit_bias = []
                mlp_l1_exit_weight  = []
                exit_output_weight = []
                exit_block_qkv_weight = []
                exit_block_qkv_bias = []
                exit_block_dense_weight = []
                exit_block_mlp_l0_weight = []
                exit_block_mlp_l0_bias = []
                exit_block_mlp_l1_weight = []
                for tp_rank, model in enumerate(models):
                    layer = model.language_model.encoder.layers[layer_id]
                    qkv_weight.append(layer.self_attention.query_key_value.weight.data)
                    dense_weight.append(layer.self_attention.dense.weight.data)
                    if use_exit_mlp:
                        mlp_l0_weight.append(layer.mlp.trunk.dense_h_to_4h.weight.data)
                        mlp_l1_weight.append(layer.mlp.trunk.dense_4h_to_h.weight.data)
                        mlp_l0_exit_weight.append(layer.mlp.branch.dense_h_to_4h.weight.data)
                        mlp_l1_exit_weight.append(layer.mlp.branch.dense_4h_to_h.weight.data)
                    else:
                        mlp_l0_weight.append(layer.mlp.dense_h_to_4h.weight.data)
                        mlp_l1_weight.append(layer.mlp.dense_4h_to_h.weight.data)
                    if has_early_exit and md.untie_exit_output_weights:
                        exit_output_weight.append(model.language_model.encoder.exit_output_weights[layer_num].data)
                    if md.linear_bias:
                        qkv_bias.append(layer.self_attention.query_key_value.bias.data)
                        if use_exit_mlp:
                            mlp_l0_bias.append(layer.mlp.trunk.dense_h_to_4h.bias.data)
                            mlp_l0_exit_bias.append(layer.mlp.branch.dense_h_to_4h.bias.data)
                        else:
                            mlp_l0_bias.append(layer.mlp.dense_h_to_4h.bias.data)
                    if use_exit_block:
                        exit_block_qkv_weight.append(layer.exit_block.self_attention.query_key_value.weight.data)
                        exit_block_dense_weight.append(layer.exit_block.self_attention.dense.weight.data)
                        exit_block_mlp_l0_weight.append(layer.exit_block.mlp.dense_h_to_4h.weight.data)
                        exit_block_mlp_l1_weight.append(layer.exit_block.mlp.dense_4h_to_h.weight.data)
                        if md.linear_bias:
                            exit_block_qkv_bias.append(layer.exit_block.self_attention.query_key_value.bias.data)
                            exit_block_mlp_l0_bias.append(layer.exit_block.mlp.dense_h_to_4h.bias.data)

                # Handle gated linear units
                if md.swiglu:
                    # concat all the first halves ('W's) and all the second halves ('V's)
                    for tp_rank in range(tp_size):
                        mlp_l0_weight[tp_rank] = torch.chunk(mlp_l0_weight[tp_rank], 2, dim=0)
                        if use_exit_mlp:
                            mlp_l0_exit_weight[tp_rank] = torch.chunk(mlp_l0_exit_weight[tp_rank], 2, dim=0)
                    message["mlp l0 weight W"] = torch.cat([w[0] for w in mlp_l0_weight], dim=0)
                    message["mlp l0 weight V"] = torch.cat([w[1] for w in mlp_l0_weight], dim=0)
                    if use_exit_mlp:
                        message["mlp l0 exit weight W"] = torch.cat([w[0] for w in mlp_l0_exit_weight], dim=0)
                        message["mlp l0 exit weight V"] = torch.cat([w[1] for w in mlp_l0_exit_weight], dim=0)
                else:
                    message["mlp l0 weight"] = torch.cat(mlp_l0_weight, dim=0)
                    if use_exit_mlp:
                        message["mlp l0 exit weight"] = torch.cat(mlp_l0_exit_weight, dim=0)

                # simple concat of the rest
                message["qkv weight"] = torch.cat(qkv_weight, dim=0)
                message["dense weight"] = torch.cat(dense_weight, dim=1)
                message["mlp l1 weight"] = torch.cat(mlp_l1_weight, dim=1)
                if use_exit_mlp:
                    message["mlp l1 exit weight"] = torch.cat(mlp_l1_exit_weight, dim=1)
                if has_early_exit and md.untie_exit_output_weights:
                    message["exit output weight"] = torch.cat(exit_output_weight, dim=0)
                if md.linear_bias:
                    message["qkv bias"] = torch.cat(qkv_bias, dim=0)
                    if md.swiglu:
                        for tp_rank in range(tp_size):
                            mlp_l0_bias[tp_rank] = torch.chunk(mlp_l0_bias[tp_rank], 2, dim=0)
                            if use_exit_mlp:
                                mlp_l0_exit_bias[tp_rank] = torch.chunk(mlp_l0_exit_bias[tp_rank], 2, dim=0)
                        message["mlp l0 bias W"] = torch.cat([b[0] for b in mlp_l0_bias],dim=0)
                        message["mlp l0 bias V"] = torch.cat([b[1] for b in mlp_l0_bias],dim=0)
                        if use_exit_mlp:
                            message["mlp l0 exit bias W"] = torch.cat([b[0] for b in mlp_l0_exit_bias],dim=0)
                            message["mlp l0 exit bias V"] = torch.cat([b[1] for b in mlp_l0_exit_bias],dim=0)
                    else:
                        message["mlp l0 bias"] = torch.cat(mlp_l0_bias, dim=0)
                        if use_exit_mlp:
                            message["mlp l0 exit bias"] = torch.cat(mlp_l0_exit_bias, dim=0)
                if use_exit_block:
                    if md.swiglu:
                        for tp_rank in range(tp_size):
                            exit_block_mlp_l0_weight[tp_rank] = torch.chunk(exit_block_mlp_l0_weight[tp_rank], 2, dim=0)
                        message["exit block mlp l0 weight W"] = torch.cat([w[0] for w in exit_block_mlp_l0_weight], dim=0)
                        message["exit block mlp l0 weight V"] = torch.cat([w[1] for w in exit_block_mlp_l0_weigh
# ... truncated ...


        }
        queue_put("pooler", message)

        message = {
            "dense weight": models[0].lm_head.dense.weight.data,
            "dense bias": models[0].lm_head.dense.bias.data,
            "norm weight": models[0].lm_head.norm.weight.data,
        }
        if norm_has_bias:
            message["norm bias"] = models[0].lm_head.norm.bias.data
        queue_put("lm head", message)

        if md.bert_binary_head:
            message = {
                "weight": models[0].binary_head.weight.data,
                "bias": models[0].binary_head.bias.data
            }
            queue_put("binary head", message)
    queue.put("done")

def load_checkpoint(queue, args):
    try:
        _load_checkpoint(queue, args)
    except:
        queue.put("exit")
        raise


                '--use-cpu-initialization',
                '--micro-batch-size', '1',
                '--no-load-optim',
                '--no-load-rng',
                '--no-save-optim',
                '--no-save-rng',
                '--no-initialization',
                '--load', args.load_dir,
                '--load-iteration', str(args.load_iteration)
                ]

    margs = parse_args()
    margs, checkpoint_args = load_args_from_checkpoint(margs)

    # Arguments do sanity checks on the world size, but we don't care,
    # so trick it into thinking we are plenty of processes
    margs.world_size = margs.tensor_model_parallel_size * margs.pipeline_model_parallel_size

    margs = validate_args(margs)

    def check_for_arg(arg_name, default=None):
        if getattr(margs, arg_name, None) is None:
            if default is not None:
                setattr(margs, arg_name, default)
            else:
                print(f"Checkpoint does not specify the argument {arg_name}. Exiting.")
                print(f"Arguments: {margs}")
                queue.put("exit")
                exit(1)

    check_for_arg('tensor_model_parallel_size')
    check_for_arg('pipeline_model_parallel_size')
    check_for_arg('num_layers')
    check_for_arg('hidden_size')
    check_for_arg('seq_length')
    check_for_arg('num_attention_heads')
    check_for_arg('max_position_embeddings')
    check_for_arg('position_embedding_type')
    check_for_arg('tokenizer_type')
    check_for_arg('iteration')
    check_for_arg('bert_binary_head')
    check_for_arg('disable_bias_linear', False)
    check_for_arg('params_dtype')
    check_for_arg('swiglu', False)

    # Determine how to make our models
    if args.model_type == 'GPT':
        from pretrain_gpt import model_provider
        margs.model_type = ModelType.encoder_or_decoder


    check_for_arg('swiglu', False)

    # Determine how to make our models
    if args.model_type == 'GPT':
        from pretrain_gpt import model_provider
        margs.model_type = ModelType.encoder_or_decoder
    elif args.model_type == 'BERT':
        from pretrain_bert import model_provider
        margs.model_type = ModelType.encoder_or_decoder
    elif args.model_type == 'EarlyExitGPT':
        from pretrain_early_exit_gpt import model_provider
        margs.model_type = ModelType.encoder_or_decoder
    else:
        raise Exception(f'unrecognized model type: {args.model_type}')

    # supress warning about torch.distributed not being initialized
    module.MegatronModule.embedding_warning_printed = True

    consumed_train_samples = None
    consumed_valid_samples = None
    def get_models(count, dtype):
        nonlocal consumed_train_samples
        nonlocal consumed_valid_samples
        model_array_len = margs.virtual_pipeline_model_parallel_size
        if model_array_len is None:
            model_array_len = 1
        models = [[] for _ in range(model_array_len)]
        pre_process = mpu.is_pipeline_first_stage()
        post_process = mpu.is_pipeline_last_stage()
        for rank in range(count):
            mpu.set_tensor_model_parallel_rank(rank)
            if margs.virtual_pipeline_model_parallel_size is not None:
                model_ = []
                for i in range(margs.virtual_pipeline_model_parallel_size):
                    mpu.set_virtual_pipeline_model_parallel_rank(i)
                    # Set pre_process and post_process only after virtual rank is set.
                    pre_process = mpu.is_pipeline_first_stage()
                    post_process = mpu.is_pipeline_last_stage()
                    this_model = model_provider(
                        pre_process=pre_process,
                        post_process=post_process
                    ).to(dtype)
                    model_.append(this_model)
            else:
                pre_process = mpu.is_pipeline_first_stage()
                post_process = mpu.is_pipeline_last_stage()
                model_rank = 0
                model_ = [model_provider(pre_process, post_process).to(dtype)]
            margs.consumed_train_samples = 0
            margs.consumed_valid_samples = 0
            load_checkpoint(model_, None, None)

            if consumed_train_samples is not None:
                assert(margs.consumed_train_samples == consumed_train_samples)
            else:
                consumed_train_samples = margs.consumed_train_samples
            if consumed_valid_samples is not None:
                assert(margs.consumed_valid_samples == consumed_valid_samples)
            else:
                consumed_valid_samples = margs.consumed_valid_samples
            for vp_rank in range(model_array_len):
                models[vp_rank].append(model_[vp_rank])
        return models

    set_global_variables(margs, build_tokenizer=False)
    mpu.set_tensor_model_parallel_world_size(margs.tensor_model_parallel_size)
    mpu.set_pipeline_model_parallel_world_size(margs.pipeline_model_parallel_size)
    mpu.set_virtual_pipeline_model_parallel_world_size(margs.virtual_pipeline_model_parallel_size)
    fused_kernels.load(margs)

    # Get true (non-padded) vocab size
    if args.true_vocab_size is not None:
        true_vocab_size = args.true_vocab_size
    elif args.vocab_file is not None:
        vocab = json.load(open(args.vocab_file))
        true_vocab_size = len(vocab)
        if args.true_vocab_size is not None and true_vocab_size != args.true_vocab_size:
            print("Both --true-vocab-size and --vocab-file specified and the vocab size does not match, aborting.")
            queue.put("exit")
            exit(1)
    else:
        true_vocab_size = None


    md.use_exit_mlp = margs.use_exit_mlp if hasattr(margs, 'use_exit_mlp') else False
    md.use_exit_block = margs.use_exit_block if hasattr(margs, 'use_exit_block') else False
    md.use_exit_norm = margs.use_exit_norm if hasattr(margs, 'use_exit_norm') else False
    md.untie_exit_output_weights = margs.untie_exit_output_weights if hasattr(margs, 'untie_exit_output_weights') else False
    md.pre_exit = margs.pre_exit
    md.checkpoint_args = checkpoint_args

    # Get first pipe stage
    mpu.set_pipeline_model_parallel_rank(0)
    if len(md.exit_layer_nums) > 0:
        layer_per_stage = md.num_layers / margs.pipeline_model_parallel_size
        mpu.set_early_exit_layer_nums(list(filter(lambda x: 0 < x <= layer_per_stage, md.exit_layer_nums)))
        mpu.set_early_exit_stages(list(set(map(lambda layer_num:  int((layer_num - 1) // layer_per_stage), md.exit_layer_nums))))
    all_models = [get_models(tp_size, md.params_dtype)]
    models = all_models[0][0]

    md.consumed_train_samples = consumed_train_samples
    md.consumed_valid_samples = consumed_valid_samples
    queue.put(md)

    def queue_put(name, msg):
        print(f"sending {name}")
        msg["name"] = name
        queue.put(msg)

    # Send embeddings
    message = {
        "word embeddings": torch.cat(
            [models[tp_rank].language_model.embedding.word_embeddings.weight.data for tp_rank in range(tp_size)],
            dim = 0)
    }
    if md.position_embedding_type == 'learned_absolute':
        message["position embeddings"] = models[0].language_model.embedding.position_embeddings.weight.data
    else:
        assert not hasattr(models[0].language_model.embedding, 'position_embeddings')

    queue_put("embeddings", message)

    total_layer_num = 0
    for vp_rank in range(vp_size):
        mpu.set_virtual_pipeline_model_parallel_rank(vp_rank)
        for pp_rank in range(pp_size):
            if pp_rank > 0:
                mpu.set_pipeline_model_parallel_rank(pp_rank)


import json
import os
import sys
import torch
import argparse
import math
from collections import OrderedDict

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--load-dir', type=str)
    parser.add_argument('--load-iteration', type=int)
    parser.add_argument('--save-dir', type=str)
    parser.add_argument('--conversion-type', choices=['exit-position', 'add-exit'], default='add-exit')
    parser.add_argument('--target-exit-position', choices=['pre', 'post'], default='post')
    parser.add_argument('--add-exit-layer-nums', type=int, nargs='+', default=[])
    parser.add_argument('--use-exit-mlp', action='store_true')
    parser.add_argument('--use-exit-block', action='store_true')
    parser.add_argument('--use-exit-norm', action='store_true')
    parser.add_argument('--random-init', action='store_true')
    parser.add_argument('--init-method-std', type=float, default=0.02)
    parser.add_argument('--megatron-path', type=str, default=None)
    return parser.parse_args()

def load_checkpoint_args(checkpoint_root_path):
    if os.path.exists(os.path.join(checkpoint_root_path, 'mp_rank_00')):
        checkpoint_rank_0_dir = 'mp_rank_00'
    elif os.path.exists(os.path.join(checkpoint_root_path, 'mp_rank_00_000')):
        checkpoint_rank_0_dir = 'mp_rank_00_000'
    else:
        raise FileNotFoundError(f'Checkpoint file {checkpoint_root_path} not found')
    checkpoint_path = os.path.join(checkpoint_root_path, checkpoint_rank_0_dir, 'model_optim_rng.pt')
    print(f"Loading args from {checkpoint_root_path}")
    model = torch.load(checkpoint_path)
    return model['args']

# Init method from megatron-lm
def init_method_normal(sigma):

    def init_(tensor):
        return torch.nn.init.normal_(tensor, mean=0.0, std=sigma)

    return init_

def scaled_init_method_normal(sigma, num_layers):
    std = sigma / math.sqrt(2.0 * num_layers)

    def init_(tensor):
        return torch.nn.init.normal_(tensor, mean=0.0, std=std)

    return init_

def change_exit_position(args, checkpoint_load_dir, checkpoint_save_dir):
    checkpoint_args = load_checkpoint_args(checkpoint_load_dir)
    cur_exit_position = 'pre' if checkpoint_args.pre_exit else 'post'
    if cur_exit_position == args.target_exit_position:
        print("No need to convert")
        return
    pipeline_parallel_size = checkpoint_args.pipeline_model_parallel_size
    tensor_parallel_size = checkpoint_args.tensor_model_parallel_size
    exit_layer_nums = checkpoint_args.exit_layer_nums
    if args.target_exit_position == 'pre':
        exit_layer_nums = [layer_num + 1 for layer_num in exit_layer_nums]
    else:
        exit_layer_nums = [layer_num - 1 for layer_num in exit_layer_nums]
    use_pipeline_parallel = pipeline_parallel_size > 1
    for tensor_rank in range(tensor_parallel_size):
        checkpoint_dicts = {}
        exit_output_weights = []
        exit_output_weight_offset = 0
        # load all pipeline ranks
        for pipeline_rank in range(pipeline_parallel_size):
            if not use_pipeline_parallel:
                checkpoint_name = os.path.join(checkpoint_load_dir, f'mp_rank_{tensor_rank:02d}', 'model_optim_rng.pt')
            else:
                checkpoint_name = os.path.join(checkpoint_load_dir, f'mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}', 'model_optim_rng.pt')
            print(f'Loading checkpoint [pp:{pipeline_rank}, tp:{tensor_rank}] from {checkpoint_name} ...')
            state_dict = torch.load(checkpoint_name, map_location='cpu')
            checkpoint_dicts[pipeline_rank] = state_dict
            # convert args
            state_dict['args'].exit_layer_nums = exit_layer_nums
            state_dict['args'].pre_exit = (args.target_exit_position == 'pre')
            # get exit output weight
            if checkpoint_args.untie_exit_output_weights and use_pipeline_parallel:
                if 'exit_output_layer' in state_dict['model']['language_model']:
                    exit_weight_num = len(state_dict['model']['language_model']['exit_output_layer'])
                    for i in range(exit_weight_num):
                        exit_output_weights.append(state_dict['model']['language_model']['exit_output_layer'].pop(f'{i}.weight'))
        # convert output weight position
        if checkpoint_args.untie_exit_output_weights and use_pipeline_parallel:
            layer_per_stage = checkpoint_args.num_layers / pipeline_parallel_size
            for pipeline_rank in range(pipeline_parallel_size):
                layer_nums = list(filter(lambda x: (layer_per_stage * pipeline_rank + 1) <= x <= (layer_per_stage * (pipeline_rank + 1)), exit_layer_nums))
                if len(layer_nums) > 0:
                    if 'exit_output_layer' not in checkpoint_dicts[pipeline_rank]['model']['language_model']:
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['exit_output_layer'] = OrderedDict()
                    for i in range(len(layer_nums)):
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['exit_output_layer'][f'{i}.weight'] = exit_output_weights[exit_output_weight_offset]
                        exit_output_weight_offset += 1
                elif 'exit_output_layer' in checkpoint_dicts[pipeline_rank]['model']['language_model']:
                    checkpoint_dicts[pipeline_rank]['model']['language_model'].pop('exit_output_layer')
        # save back
        for pipeline_rank in range(pipeline_parallel_size):
            if not use_pipeline_parallel:
                checkpoint_save_path = os.path.join(checkpoint_save_dir, f'mp_rank_{tensor_rank:02d}', 'model_optim_rng.pt')
            else:
                checkpoint_save_path = os.path.join(checkpoint_save_dir, f'mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}', 'model_optim_rng.pt')
            dirname = os.path.dirname(checkpoint_save_path)
            os.makedirs(dirname, exist_ok = True)
            print(f'Saving checkpoint [pp:{pipeline_rank}, tp:{tensor_rank}] to {checkpoint_save_path} ...')
            torch.save(checkpoint_dicts[pipeline_rank], checkpoint_save_path)
    print('Exit Weight Position Conversion Completed')

def add_exit(args, checkpoint_load_dir, checkpoint_save_dir):
    if len(args.add_exit_layer_nums) == 0:
        print("No exit layer to add")
        return
    checkpoint_args = load_checkpoint_args(checkpoint_load_dir)
    use_pre_exit = False
    if len(checkpoint_args.exit_layer_nums) == 0:
        if args.target_exit_position == 'pre':
            use_pre_exit = True
    else:
        if checkpoint_args.pre_exit == (args.target_exit_position == 'pre'):
            print("Can't add exit layers and change exit position at the same time")
            return
        use_pre_exit = checkpoint_args.pre_exit
    target_exit_layer_nums = sorted(list(set(checkpoint_args.exit_layer_nums + args.add_exit_layer_nums)))
    tensor_parallel_size = checkpoint_args.tensor_model_parallel_size
    pipeline_parallel_size = checkpoint_args.pipeline_model_parallel_size
    use_pipeline_parallel = pipeline_parallel_size > 1
    layer_per_stage = checkpoint_args.num_layers / pipeline_parallel_size
    # if args.random_init:
    #     init_method = init_method_normal(args.init_method_std)
    #     output_layer_init_method = scaled_init_method_normal(args.init_method_std)

    for tensor_rank in range(tensor_parallel_size):
        checkpoint_dicts = {}
        output_weight = None
        final_norm_weight = None
        final_norm_bias = None
        # load all pipeline ranks
        for pipeline_rank in range(pipeline_parallel_size):
            if not use_pipeline_parallel:
                checkpoint_name = os.path.join(checkpoint_load_dir, f'mp_rank_{tensor_rank:02d}', 'model_optim_rng.pt')
            else:
                checkpoint_name = os.path.join(checkpoint_load_dir, f'mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}', 'model_optim_rng.pt')
            print(f'Loading checkpoint [pp:{pipeline_rank}, tp:{tensor_rank}] from {checkpoint_name} ...')
            layer_num_offset = layer_per_stage * pipeline_rank + 1
            exit_layer_nums = list(filter(lambda x: (layer_per_stage * pipeline_rank + 1) <= x <= (layer_per_stage * (pipeline_rank + 1)), target_exit_layer_nums))
            state_dict = torch.load(checkpoint_name)
            checkpoint_dicts[pipeline_rank] = state_dict
            # convert args
            state_dict['args'].exit_layer_nums = target_exit_layer_nums
            state_dict['args'].pre_exit = use_pre_exit
            state_dict['args'].untie_exit_output_weights = True

            # get ouptut weight
            if checkpoint_args.untie_embeddings_and_output_weights:
                if pipeline_rank == pipeline_parallel_size - 1:
                    output_weight = state_dict['model']['language_model']['output_layer']['weight']
            else:
                if pipeline_rank == 0:
                    output_weight = state_dict['model']['language_model']['embedding']['word_embeddings']['weight']

            # convert to exit mlp
            if args.use_exit_mlp and (not hasattr(state_dict['args'], 'use_exit_mlp') or not state_dict['args'].use_exit_mlp):
                state_dict['args'].use_exit_mlp = args.use_exit_mlp
                for layer_num in exit_layer_nums:
                    if args.random_init:
                        init_method = init_method_normal(args.init_method_std)
                        output_layer_init_method = scaled_init_method_normal(args.init_method_std, layer_num)
                    layer_id = int(layer_num - layer_num_offset)
                    state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.trunk.dense_h_to_4h.weight'] = \
                            state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_h_to_4h.weight']
                    state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.trunk.dense_4h_to_h.weight'] = \
                            state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_4h_to_h.weight']
                    if args.random_init:
                        state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.branch.dense_h_to_4h.weight'] = \
                                init_method(torch.empty(state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_h_to_4h.weight'].shape))
                        state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.branch.dense_4h_to_h.weight'] = \
                                output_layer_init_method(torch.empty(state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_4h_to_h.weight'].shape))
                    else:
                        state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.branch.dense_h_to_4h.weight'] = \
                                state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_h_to_4h.weight']
                        state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.branch.dense_4h_to_h.weight'] = \
                                state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_4h_to_h.weight']
                    state_dict['model']['language_model']['encoder'].pop(f'layers.{layer_id}.mlp.dense_h_to_4h.weight')
                    state_dict['model']['language_model']['encoder'].pop(f'layers.{layer_id}.mlp.dense_4h_to_h.weight')
                    if checkpoint_args.add_bias_linear:
                        state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.trunk.dense_h_to_4h.bias'] = \
                                state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_h_to_4h.bias']
                        state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.trunk.dense_4h_to_h.bias'] = \
                                state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_4h_to_h.bias']
                        if args.random_init:
                            state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.branch.dense_h_to_4h.bias'] = \
                                    torch.zeros(state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_h_to_4h.bias'].shape)
                            state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.branch.dense_4h_to_h.bias'] = \
                                    torch.zeros(state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_4h_to_h.bias'].shape)
                        else:
                            state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.branch.dense_h_to_4h.bias'] = \
                                    state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_h_to_4h.bias']
                            state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.branch.dense_4h_to_h.bias'] = \
                                    state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_4h_to_h.bias']
                        state_dict['model']['language_model']['encoder'].pop(f'layers.{layer_id}.mlp.dense_h_to_4h.bias')
                        state_dict['model']['language_model']['encoder'].pop(f'layers.{layer_id}.mlp.dense_4h_to_h.bias')
            # convert to exit block
            if args.use_exit_block:
                state_dict['args'].use_exit_block = args.use_exit_block
                # get last layer params
                if pipeline_rank == pipeline_parallel_size - 1:
                    last_layer_id = int(layer_per_stage - 1)
                    last_layer_input_norm = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.input_norm.weight']
                    last_layer_atten_qkv = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.self_attention.query_key_value.weight']
                    last_layer_atten_dense = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.self_attention.dense.weight']
                    last_layer_post_norm = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.post_attention_norm.weight']
                    last_layer_mlp_h_to_4h = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.mlp.dense_h_to_4h.weight']
                    last_layer_mlp_4h_to_h = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.mlp.dense_4h_to_h.weight']
                    if checkpoint_args.add_bias_linear:
                        last_layer_atten_dense_bias = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.self_attention.dense.bias']
                        last_layer_h_to_4h_bias = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.mlp.dense_h_to_4h.bias']
                        last_layer_4h_to_h_bias = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.mlp.dense_4h_to_h.bias']
                    if checkpoint_args.normalization == 'LayerNorm':
                        last_layer_input_norm_bias = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.input_norm.bias']
                        last_layer_post_norm_bias = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.post_attention_norm.bias']
            # get final norm
            if args.use_exit_norm:
                state_dict['args'].use_exit_norm = args.use_exit_norm
                if 'final_norm.weight' in state_dict['model']['language_model']['encoder']:
                    final_norm_weight = state_dict['model']['language_model']['encoder']['final_norm.weight']
                    if checkpoint_args.normalization == 'LayerNorm':
                        final_norm_bias = state_dict['model']['language_mode']['encoder']['final_norm.bias']
            # get exit output weight
            if len(exit_layer_nums) > 0 and 'exit_output_layer' not in state_dict['model']['language_model']:
                state_dict['model']['language_model']['exit_output_layer'] = OrderedDict()

        for pipeline_rank in range(pipeline_parallel_size):
            layer_num_offset = layer_per_stage * pipeline_rank + 1
            exit_layer_nums = list(filter(lambda x: (layer_per_stage * pipeline_rank + 1) <= x <= (layer_per_stage * (pipeline_rank + 1)), target_exit_layer_nums))
            # add exit output weight and exit norm
            for i, layer_num in enumerate(exit_layer_nums):
                layer_id = int(layer_num - layer_num_offset)
                if args.random_init:
                    init_method = init_method_normal(args.init_method_std)
                    output_layer_init_method = scaled_init_method_normal(args.init_method_std, layer_num)
                    checkpoint_dicts[pipeline_rank]['model']['language_model']['exit_output_layer'][f'{i}.weight'] = init_method(torch.empty(output_weight.shape))
                else:
                    checkpoint_dicts[pipeline_rank]['model']['language_model']['exit_output_layer'][f'{i}.weight'] = output_weight
                if args.use_exit_block:
                    if args.random_init:
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.input_norm.weight'] = torch.ones(last_layer_input_norm.shape)
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.self_attention.query_key_value.weight'] = init_method(torch.empty(last_layer_atten_qkv.shape))
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.self_attention.dense.weight'] =output_layer_init_method(torch.empty(last_layer_atten_dense.shape))
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.post_attention_norm.weight'] = torch.ones(last_layer_post_norm.shape)
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.mlp.dense_h_to_4h.weight'] = init_method(torch.empty(last_layer_mlp_h_to_4h.shape))
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.mlp.dense_4h_to_h.weight'] = output_layer_init_method(torch.empty(last_layer_mlp_4h_to_h.shape))
                        if checkpoint_args.add_bias_linear:
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.self_attention.dense.bias'] = torch.zeros(last_layer_atten_dense_bias.shape)
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.mlp.dense_h_to_4h.bias'] = torch.zeros(last_layer_h_to_4h_bias.shape)
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.mlp.dense_4h_to_h.bias'] = torch.zeros(last_layer_4h_to_h_bias.shape)
                        if checkpoint_args.normalization == 'LayerNorm':
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.input_norm.bias'] = torch.zeros(last_layer_input_norm_bias.shape)
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.post_attention_norm.bias'] = torch.zeros(last_layer_post_norm_bias.shape)
                    else:
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.input_norm.weight'] = last_layer_input_
# ... truncated ...


import json
import os
import sys
import torch
import argparse
import math
from collections import OrderedDict

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--load-dir', type=str)
    parser.add_argument('--load-iteration', type=int)
    parser.add_argument('--save-dir', type=str)
    parser.add_argument('--conversion-type', choices=['exit-position', 'add-exit'], default='add-exit')
    parser.add_argument('--target-exit-position', choices=['pre', 'post'], default='post')
    parser.add_argument('--add-exit-layer-nums', type=int, nargs='+', default=[])
    parser.add_argument('--use-exit-mlp', action='store_true')
    parser.add_argument('--use-exit-block', action='store_true')
    parser.add_argument('--use-exit-norm', action='store_true')
    parser.add_argument('--random-init', action='store_true')
    parser.add_argument('--init-method-std', type=float, default=0.02)
    parser.add_argument('--megatron-path', type=str, default=None)
    return parser.parse_args()

def load_checkpoint_args(checkpoint_root_path):
    if os.path.exists(os.path.join(checkpoint_root_path, 'mp_rank_00')):
        checkpoint_rank_0_dir = 'mp_rank_00'
    elif os.path.exists(os.path.join(checkpoint_root_path, 'mp_rank_00_000')):
        checkpoint_rank_0_dir = 'mp_rank_00_000'
    else:
        raise FileNotFoundError(f'Checkpoint file {checkpoint_root_path} not found')
    checkpoint_path = os.path.join(checkpoint_root_path, checkpoint_rank_0_dir, 'model_optim_rng.pt')
    print(f"Loading args from {checkpoint_root_path}")
    model = torch.load(checkpoint_path)
    return model['args']

# Init method from megatron-lm
def init_method_normal(sigma):

    def init_(tensor):
        return torch.nn.init.normal_(tensor, mean=0.0, std=sigma)

    return init_


import argparse
import math
from collections import OrderedDict

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--load-dir', type=str)
    parser.add_argument('--load-iteration', type=int)
    parser.add_argument('--save-dir', type=str)
    parser.add_argument('--conversion-type', choices=['exit-position', 'add-exit'], default='add-exit')
    parser.add_argument('--target-exit-position', choices=['pre', 'post'], default='post')
    parser.add_argument('--add-exit-layer-nums', type=int, nargs='+', default=[])
    parser.add_argument('--use-exit-mlp', action='store_true')
    parser.add_argument('--use-exit-block', action='store_true')
    parser.add_argument('--use-exit-norm', action='store_true')
    parser.add_argument('--random-init', action='store_true')
    parser.add_argument('--init-method-std', type=float, default=0.02)
    parser.add_argument('--megatron-path', type=str, default=None)
    return parser.parse_args()

def load_checkpoint_args(checkpoint_root_path):
    if os.path.exists(os.path.join(checkpoint_root_path, 'mp_rank_00')):
        checkpoint_rank_0_dir = 'mp_rank_00'
    elif os.path.exists(os.path.join(checkpoint_root_path, 'mp_rank_00_000')):
        checkpoint_rank_0_dir = 'mp_rank_00_000'
    else:
        raise FileNotFoundError(f'Checkpoint file {checkpoint_root_path} not found')
    checkpoint_path = os.path.join(checkpoint_root_path, checkpoint_rank_0_dir, 'model_optim_rng.pt')
    print(f"Loading args from {checkpoint_root_path}")
    model = torch.load(checkpoint_path)
    return model['args']

# Init method from megatron-lm
def init_method_normal(sigma):

    def init_(tensor):
        return torch.nn.init.normal_(tensor, mean=0.0, std=sigma)

    return init_

def scaled_init_method_normal(sigma, num_layers):
    std = sigma / math.sqrt(2.0 * num_layers)

    def init_(tensor):
        return torch.nn.init.normal_(tensor, mean=0.0, std=std)

    return init_

def change_exit_position(args, checkpoint_load_dir, checkpoint_save_dir):
    checkpoint_args = load_checkpoint_args(checkpoint_load_dir)


    parser.add_argument('--use-exit-block', action='store_true')
    parser.add_argument('--use-exit-norm', action='store_true')
    parser.add_argument('--random-init', action='store_true')
    parser.add_argument('--init-method-std', type=float, default=0.02)
    parser.add_argument('--megatron-path', type=str, default=None)
    return parser.parse_args()

def load_checkpoint_args(checkpoint_root_path):
    if os.path.exists(os.path.join(checkpoint_root_path, 'mp_rank_00')):
        checkpoint_rank_0_dir = 'mp_rank_00'
    elif os.path.exists(os.path.join(checkpoint_root_path, 'mp_rank_00_000')):
        checkpoint_rank_0_dir = 'mp_rank_00_000'
    else:
        raise FileNotFoundError(f'Checkpoint file {checkpoint_root_path} not found')
    checkpoint_path = os.path.join(checkpoint_root_path, checkpoint_rank_0_dir, 'model_optim_rng.pt')
    print(f"Loading args from {checkpoint_root_path}")
    model = torch.load(checkpoint_path)
    return model['args']

# Init method from megatron-lm
def init_method_normal(sigma):

    def init_(tensor):
        return torch.nn.init.normal_(tensor, mean=0.0, std=sigma)

    return init_

def scaled_init_method_normal(sigma, num_layers):
    std = sigma / math.sqrt(2.0 * num_layers)

    def init_(tensor):
        return torch.nn.init.normal_(tensor, mean=0.0, std=std)

    return init_

def change_exit_position(args, checkpoint_load_dir, checkpoint_save_dir):
    checkpoint_args = load_checkpoint_args(checkpoint_load_dir)
    cur_exit_position = 'pre' if checkpoint_args.pre_exit else 'post'
    if cur_exit_position == args.target_exit_position:
        print("No need to convert")
        return
    pipeline_parallel_size = checkpoint_args.pipeline_model_parallel_size
    tensor_parallel_size = checkpoint_args.tensor_model_parallel_size
    exit_layer_nums = checkpoint_args.exit_layer_nums
    if args.target_exit_position == 'pre':


def load_checkpoint_args(checkpoint_root_path):
    if os.path.exists(os.path.join(checkpoint_root_path, 'mp_rank_00')):
        checkpoint_rank_0_dir = 'mp_rank_00'
    elif os.path.exists(os.path.join(checkpoint_root_path, 'mp_rank_00_000')):
        checkpoint_rank_0_dir = 'mp_rank_00_000'
    else:
        raise FileNotFoundError(f'Checkpoint file {checkpoint_root_path} not found')
    checkpoint_path = os.path.join(checkpoint_root_path, checkpoint_rank_0_dir, 'model_optim_rng.pt')
    print(f"Loading args from {checkpoint_root_path}")
    model = torch.load(checkpoint_path)
    return model['args']

# Init method from megatron-lm
def init_method_normal(sigma):

    def init_(tensor):
        return torch.nn.init.normal_(tensor, mean=0.0, std=sigma)

    return init_

def scaled_init_method_normal(sigma, num_layers):
    std = sigma / math.sqrt(2.0 * num_layers)

    def init_(tensor):
        return torch.nn.init.normal_(tensor, mean=0.0, std=std)

    return init_

def change_exit_position(args, checkpoint_load_dir, checkpoint_save_dir):
    checkpoint_args = load_checkpoint_args(checkpoint_load_dir)
    cur_exit_position = 'pre' if checkpoint_args.pre_exit else 'post'
    if cur_exit_position == args.target_exit_position:
        print("No need to convert")
        return
    pipeline_parallel_size = checkpoint_args.pipeline_model_parallel_size
    tensor_parallel_size = checkpoint_args.tensor_model_parallel_size
    exit_layer_nums = checkpoint_args.exit_layer_nums
    if args.target_exit_position == 'pre':
        exit_layer_nums = [layer_num + 1 for layer_num in exit_layer_nums]
    else:
        exit_layer_nums = [layer_num - 1 for layer_num in exit_layer_nums]
    use_pipeline_parallel = pipeline_parallel_size > 1
    for tensor_rank in range(tensor_parallel_size):
        checkpoint_dicts = {}
        exit_output_weights = []
        exit_output_weight_offset = 0


    model = torch.load(checkpoint_path)
    return model['args']

# Init method from megatron-lm
def init_method_normal(sigma):

    def init_(tensor):
        return torch.nn.init.normal_(tensor, mean=0.0, std=sigma)

    return init_

def scaled_init_method_normal(sigma, num_layers):
    std = sigma / math.sqrt(2.0 * num_layers)

    def init_(tensor):
        return torch.nn.init.normal_(tensor, mean=0.0, std=std)

    return init_

def change_exit_position(args, checkpoint_load_dir, checkpoint_save_dir):
    checkpoint_args = load_checkpoint_args(checkpoint_load_dir)
    cur_exit_position = 'pre' if checkpoint_args.pre_exit else 'post'
    if cur_exit_position == args.target_exit_position:
        print("No need to convert")
        return
    pipeline_parallel_size = checkpoint_args.pipeline_model_parallel_size
    tensor_parallel_size = checkpoint_args.tensor_model_parallel_size
    exit_layer_nums = checkpoint_args.exit_layer_nums
    if args.target_exit_position == 'pre':
        exit_layer_nums = [layer_num + 1 for layer_num in exit_layer_nums]
    else:
        exit_layer_nums = [layer_num - 1 for layer_num in exit_layer_nums]
    use_pipeline_parallel = pipeline_parallel_size > 1
    for tensor_rank in range(tensor_parallel_size):
        checkpoint_dicts = {}
        exit_output_weights = []
        exit_output_weight_offset = 0
        # load all pipeline ranks
        for pipeline_rank in range(pipeline_parallel_size):
            if not use_pipeline_parallel:
                checkpoint_name = os.path.join(checkpoint_load_dir, f'mp_rank_{tensor_rank:02d}', 'model_optim_rng.pt')
            else:
                checkpoint_name = os.path.join(checkpoint_load_dir, f'mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}', 'model_optim_rng.pt')
            print(f'Loading checkpoint [pp:{pipeline_rank}, tp:{tensor_rank}] from {checkpoint_name} ...')
            state_dict = torch.load(checkpoint_name, map_location='cpu')
            checkpoint_dicts[pipeline_rank] = state_dict
            # convert args
            state_dict['args'].exit_layer_nums = exit_layer_nums
            state_dict['args'].pre_exit = (args.target_exit_position == 'pre')
            # get exit output weight
            if checkpoint_args.untie_exit_output_weights and use_pipeline_parallel:
                if 'exit_output_layer' in state_dict['model']['language_model']:
                    exit_weight_num = len(state_dict['model']['language_model']['exit_output_layer'])
                    for i in range(exit_weight_num):
                        exit_output_weights.append(state_dict['model']['language_model']['exit_output_layer'].pop(f'{i}.weight'))
        # convert output weight position
        if checkpoint_args.untie_exit_output_weights and use_pipeline_parallel:
            layer_per_stage = checkpoint_args.num_layers / pipeline_parallel_size
            for pipeline_rank in range(pipeline_parallel_size):
                layer_nums = list(filter(lambda x: (layer_per_stage * pipeline_rank + 1) <= x <= (layer_per_stage * (pipeline_rank + 1)), exit_layer_nums))
                if len(layer_nums) > 0:
                    if 'exit_output_layer' not in checkpoint_dicts[pipeline_rank]['model']['language_model']:
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['exit_output_layer'] = OrderedDict()
                    for i in range(len(layer_nums)):
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['exit_output_layer'][f'{i}.weight'] = exit_output_weights[exit_output_weight_offset]
                        exit_output_weight_offset += 1
                elif 'exit_output_layer' in checkpoint_dicts[pipeline_rank]['model']['language_model']:
                    checkpoint_dicts[pipeline_rank]['model']['language_model'].pop('exit_output_layer')
        # save back
        for pipeline_rank in range(pipeline_parallel_size):
            if not use_pipeline_parallel:
                checkpoint_save_path = os.path.join(checkpoint_save_dir, f'mp_rank_{tensor_rank:02d}', 'model_optim_rng.pt')
            else:
                checkpoint_save_path = os.path.join(checkpoint_save_dir, f'mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}', 'model_optim_rng.pt')
            dirname = os.path.dirname(checkpoint_save_path)
            os.makedirs(dirname, exist_ok = True)
            print(f'Saving checkpoint [pp:{pipeline_rank}, tp:{tensor_rank}] to {checkpoint_save_path} ...')
            torch.save(checkpoint_dicts[pipeline_rank], checkpoint_save_path)
    print('Exit Weight Position Conversion Completed')

def add_exit(args, checkpoint_load_dir, checkpoint_save_dir):
    if len(args.add_exit_layer_nums) == 0:
        print("No exit layer to add")
        return
    checkpoint_args = load_checkpoint_args(checkpoint_load_dir)
    use_pre_exit = False
    if len(checkpoint_args.exit_layer_nums) == 0:
        if args.target_exit_position == 'pre':
            use_pre_exit = True
    else:
        if checkpoint_args.pre_exit == (args.target_exit_position == 'pre'):
            print("Can't add exit layers and change exit position at the same time")
            return
        use_pre_exit = checkpoint_args.pre_exit
    target_exit_layer_nums = sorted(list(set(checkpoint_args.exit_layer_nums + args.add_exit_layer_nums)))
    tensor_parallel_size = checkpoint_args.tensor_model_parallel_size
    pipeline_parallel_size = checkpoint_args.pipeline_model_parallel_size
    use_pipeline_parallel = pipeline_parallel_size > 1


                    if 'exit_output_layer' not in checkpoint_dicts[pipeline_rank]['model']['language_model']:
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['exit_output_layer'] = OrderedDict()
                    for i in range(len(layer_nums)):
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['exit_output_layer'][f'{i}.weight'] = exit_output_weights[exit_output_weight_offset]
                        exit_output_weight_offset += 1
                elif 'exit_output_layer' in checkpoint_dicts[pipeline_rank]['model']['language_model']:
                    checkpoint_dicts[pipeline_rank]['model']['language_model'].pop('exit_output_layer')
        # save back
        for pipeline_rank in range(pipeline_parallel_size):
            if not use_pipeline_parallel:
                checkpoint_save_path = os.path.join(checkpoint_save_dir, f'mp_rank_{tensor_rank:02d}', 'model_optim_rng.pt')
            else:
                checkpoint_save_path = os.path.join(checkpoint_save_dir, f'mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}', 'model_optim_rng.pt')
            dirname = os.path.dirname(checkpoint_save_path)
            os.makedirs(dirname, exist_ok = True)
            print(f'Saving checkpoint [pp:{pipeline_rank}, tp:{tensor_rank}] to {checkpoint_save_path} ...')
            torch.save(checkpoint_dicts[pipeline_rank], checkpoint_save_path)
    print('Exit Weight Position Conversion Completed')

def add_exit(args, checkpoint_load_dir, checkpoint_save_dir):
    if len(args.add_exit_layer_nums) == 0:
        print("No exit layer to add")
        return
    checkpoint_args = load_checkpoint_args(checkpoint_load_dir)
    use_pre_exit = False
    if len(checkpoint_args.exit_layer_nums) == 0:
        if args.target_exit_position == 'pre':
            use_pre_exit = True
    else:
        if checkpoint_args.pre_exit == (args.target_exit_position == 'pre'):
            print("Can't add exit layers and change exit position at the same time")
            return
        use_pre_exit = checkpoint_args.pre_exit
    target_exit_layer_nums = sorted(list(set(checkpoint_args.exit_layer_nums + args.add_exit_layer_nums)))
    tensor_parallel_size = checkpoint_args.tensor_model_parallel_size
    pipeline_parallel_size = checkpoint_args.pipeline_model_parallel_size
    use_pipeline_parallel = pipeline_parallel_size > 1
    layer_per_stage = checkpoint_args.num_layers / pipeline_parallel_size
    # if args.random_init:
    #     init_method = init_method_normal(args.init_method_std)
    #     output_layer_init_method = scaled_init_method_normal(args.init_method_std)

    for tensor_rank in range(tensor_parallel_size):
        checkpoint_dicts = {}
        output_weight = None
        final_norm_weight = None
        final_norm_bias = None
        # load all pipeline ranks
        for pipeline_rank in range(pipeline_parallel_size):
            if not use_pipeline_parallel:
                checkpoint_name = os.path.join(checkpoint_load_dir, f'mp_rank_{tensor_rank:02d}', 'model_optim_rng.pt')
            else:
                checkpoint_name = os.path.join(checkpoint_load_dir, f'mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}', 'model_optim_rng.pt')
            print(f'Loading checkpoint [pp:{pipeline_rank}, tp:{tensor_rank}] from {checkpoint_name} ...')
            layer_num_offset = layer_per_stage * pipeline_rank + 1
            exit_layer_nums = list(filter(lambda x: (layer_per_stage * pipeline_rank + 1) <= x <= (layer_per_stage * (pipeline_rank + 1)), target_exit_layer_nums))
            state_dict = torch.load(checkpoint_name)
            checkpoint_dicts[pipeline_rank] = state_dict
            # convert args
            state_dict['args'].exit_layer_nums = target_exit_layer_nums
            state_dict['args'].pre_exit = use_pre_exit
            state_dict['args'].untie_exit_output_weights = True

            # get ouptut weight
            if checkpoint_args.untie_embeddings_and_output_weights:
                if pipeline_rank == pipeline_parallel_size - 1:
                    output_weight = state_dict['model']['language_model']['output_layer']['weight']
            else:
                if pipeline_rank == 0:
                    output_weight = state_dict['model']['language_model']['embedding']['word_embeddings']['weight']

            # convert to exit mlp
            if args.use_exit_mlp and (not hasattr(state_dict['args'], 'use_exit_mlp') or not state_dict['args'].use_exit_mlp):
                state_dict['args'].use_exit_mlp = args.use_exit_mlp
                for layer_num in exit_layer_nums:
                    if args.random_init:
                        init_method = init_method_normal(args.init_method_std)
                        output_layer_init_method = scaled_init_method_normal(args.init_method_std, layer_num)
                    layer_id = int(layer_num - layer_num_offset)
                    state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.trunk.dense_h_to_4h.weight'] = \
                            state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_h_to_4h.weight']
                    state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.trunk.dense_4h_to_h.weight'] = \
                            state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_4h_to_h.weight']
                    if args.random_init:
                        state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.branch.dense_h_to_4h.weight'] = \
                                init_method(torch.empty(state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_h_to_4h.weight'].shape))
                        state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.branch.dense_4h_to_h.weight'] = \
                                output_layer_init_method(torch.empty(state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_4h_to_h.weight'].shape))
                    else:
                        state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.branch.dense_h_to_4h.weight'] = \
                                state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_h_to_4h.weight']
                        state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.branch.dense_4h_to_h.weight'] = \
                                state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_4h_to_h.weight']
                    state_dict['model']['language_model']['encoder'].pop(f'layers.{layer_id}.mlp.dense_h_to_4h.weight')
                    state_dict['model']['language_model']['encoder'].pop(f'layers.{layer_id}.mlp.dense_4h_to_h.weight')
                    if checkpoint_args.add_bias_linear:
                        state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.trunk.dense_h_to_4h.bias'] = \
                                state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_h_to_4h.bias']
                        state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.trunk.dense_4h_to_h.bias'] = \
                                state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_4h_to_h.bias']
                        if args.random_init:
                            state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.branch.dense_h_to_4h.bias'] = \
                                    torch.zeros(state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_h_to_4h.bias'].shape)
                            state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.branch.dense_4h_to_h.bias'] = \
                                    torch.zeros(state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_4h_to_h.bias'].shape)
                        else:
                            state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.branch.dense_h_to_4h.bias'] = \
                                    state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_h_to_4h.bias']
                            state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.branch.dense_4h_to_h.bias'] = \
                                    state_dict['model']['language_model']['encoder'][f'layers.{layer_id}.mlp.dense_4h_to_h.bias']
                        state_dict['model']['language_model']['encoder'].pop(f'layers.{layer_id}.mlp.dense_h_to_4h.bias')
                        state_dict['model']['language_model']['encoder'].pop(f'layers.{layer_id}.mlp.dense_4h_to_h.bias')
            # convert to exit block
            if args.use_exit_block:
                state_dict['args'].use_exit_block = args.use_exit_block
                # get last layer params
                if pipeline_rank == pipeline_parallel_size - 1:
                    last_layer_id = int(layer_per_stage - 1)
                    last_layer_input_norm = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.input_norm.weight']
                    last_layer_atten_qkv = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.self_attention.query_key_value.weight']
                    last_layer_atten_dense = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.self_attention.dense.weight']
                    last_layer_post_norm = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.post_attention_norm.weight']
                    last_layer_mlp_h_to_4h = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.mlp.dense_h_to_4h.weight']
                    last_layer_mlp_4h_to_h = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.mlp.dense_4h_to_h.weight']
                    if checkpoint_args.add_bias_linear:
                        last_layer_atten_dense_bias = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.self_attention.dense.bias']
                        last_layer_h_to_4h_bias = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.mlp.dense_h_to_4h.bias']
                        last_layer_4h_to_h_bias = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.mlp.dense_4h_to_h.bias']
                    if checkpoint_args.normalization == 'LayerNorm':
                        last_layer_input_norm_bias = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.input_norm.bias']
                        last_layer_post_norm_bias = state_dict['model']['language_model']['encoder'][f'layers.{last_layer_id}.post_attention_norm.bias']
            # get final norm
            if args.use_exit_norm:
                state_dict['args'].use_exit_norm = args.use_exit_norm
                if 'final_norm.weight' in state_dict['model']['language_model']['encoder']:
                    final_norm_weight = state_dict['model']['language_model']['encoder']['final_norm.weight']
                    if checkpoint_args.normalization == 'LayerNorm':
                        final_norm_bias = state_dict['model']['language_mode']['encoder']['final_norm.bias']
            # get exit output weight
            if len(exit_layer_nums) > 0 and 'exit_output_layer' not in state_dict['model']['language_model']:
                state_dict['model']['language_model']['exit_output_layer'] = OrderedDict()

        for pipeline_rank in range(pipeline_parallel_size):
            layer_num_offset = layer_per_stage * pipeline_rank + 1
            exit_layer_nums = list(filter(lambda x: (layer_per_stage * pipeline_rank + 1) <= x <= (layer_per_stage * (pipeline_rank + 1)), target_exit_layer_nums))
            # add exit output weight and exit norm
            for i, layer_num in enumerate(exit_layer_nums):
                layer_id = int(layer_num - layer_num_offset)
                if args.random_init:
                    init_method = init_method_normal(args.init_method_std)
                    output_layer_init_method = scaled_init_method_normal(args.init_method_std, layer_num)
                    checkpoint_dicts[pipeline_rank]['model']['language_model']['exit_output_layer'][f'{i}.weight'] = init_method(torch.empty(output_weight.shape))
                else:
                    checkpoint_dicts[pipeline_rank]['model']['language_model']['exit_output_layer'][f'{i}.weight'] = output_weight
                if args.use_exit_block:
                    if args.random_init:
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.input_norm.weight'] = torch.ones(last_layer_input_norm.shape)
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.self_attention.query_key_value.weight'] = init_method(torch.empty(last_layer_atten_qkv.shape))
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.self_attention.dense.weight'] =output_layer_init_method(torch.empty(last_layer_atten_dense.shape))
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.post_attention_norm.weight'] = torch.ones(last_layer_post_norm.shape)
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.mlp.dense_h_to_4h.weight'] = init_method(torch.empty(last_layer_mlp_h_to_4h.shape))
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.mlp.dense_4h_to_h.weight'] = output_layer_init_method(torch.empty(last_layer_mlp_4h_to_h.shape))
                        if checkpoint_args.add_bias_linear:
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.self_attention.dense.bias'] = torch.zeros(last_layer_atten_dense_bias.shape)
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.mlp.dense_h_to_4h.bias'] = torch.zeros(last_layer_h_to_4h_bias.shape)
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.mlp.dense_4h_to_h.bias'] = torch.zeros(last_layer_4h_to_h_bias.shape)
                        if checkpoint_args.normalization == 'LayerNorm':
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.input_norm.bias'] = torch.zeros(last_layer_input_norm_bias.shape)
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.post_attention_norm.bias'] = torch.zeros(last_layer_post_norm_bias.shape)
                    else:
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.input_norm.weight'] = last_layer_input_norm
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.self_attention.query_key_value.weight'] = last_layer_atten_qkv
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.self_attention.dense.weight'] = last_layer_atten_dense
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.post_attention_norm.weight'] = last_layer_post_norm
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.mlp.dense_h_to_4h.weight'] = last_layer_mlp_h_to_4h
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.mlp.dense_4h_to_h.weight'] = last_layer_mlp_4h_to_h
                        if checkpoint_args.add_bias_linear:
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.self_attention.dense.bias'] = last_layer_atten_dense_bias
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.mlp.dense_h_to_4h.bias'] = last_layer_h_to_4h_bias
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.mlp.dense_4h_to_h.bias'] = last_layer_4h_to_h_bias
                        if checkpoint_args.normalization == 'LayerNorm':
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.input_norm.bias'] = last_layer_input_norm_bias
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_block.post_attention_norm.bias'] = last_layer_post_norm_bias
                if args.use_exit_norm:
                    if args.random_init:
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_norm.weight'] = torch.ones(final_norm_weight.shape)
                    else:
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_norm.weight'] = final_norm_weight
                    if final_norm_bias is not None:
                        if args.random_init:
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_norm.bias'] = torch.zeros(final_norm_bias.shape)
                        else:
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_norm.bias'] = final_norm_bias
            if not use_pipeline_parallel:
                checkpoint_save_path = os.path.join(checkpoint_save_dir, f'mp_rank_{tensor_rank:02d}', 'model_optim_rng.pt')
            else:
                checkpoint_save_path = os.path.join(checkpoint_save_dir, f'mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}', 'model_optim_rng.pt')
            dirname = os.path.dirname(checkpoint_save_path)
            os.makedirs(dirname, exist_ok = True)
            print(f'Saving checkpoint [pp:{pipeline_rank}, tp:{tensor_rank}] to {checkpoint_save_path} ...')
            torch.save(checkpoint_dicts[pipeline_rank], checkpoint_save_path)
    print('Add Exit Layers Completed')

def convert(args):
    if args.megatron_path is not None:
        sys.path.insert(0, args.megatron_path)
    checkpoint_load_dir = os.path.join(args.load_dir, 'iter_{:07d}'.format(args.load_iteration))
    checkpoint_save_dir = os.path.join(args.save_dir, 'iter_{:07d}'.format(args.load_iteration))
    if args.conversion_type == 'exit-position':
        change_exit_position(args, checkpoint_load_dir, checkpoint_save_dir)
    elif args.conversion_type == 'add-exit':
        add_exit(args, checkpoint_load_dir, checkpoint_save_dir)
    with open(os.path.join(args.save_dir, 'latest_checkpointed_iteration.txt'), 'w', encoding='utf-8') as f:
        f.write(str(args.load_iteration))

if __name__ == '__main__':
    args = get_args()
    convert(args)


convert

                    if args.random_init:
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_norm.weight'] = torch.ones(final_norm_weight.shape)
                    else:
                        checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_norm.weight'] = final_norm_weight
                    if final_norm_bias is not None:
                        if args.random_init:
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_norm.bias'] = torch.zeros(final_norm_bias.shape)
                        else:
                            checkpoint_dicts[pipeline_rank]['model']['language_model']['encoder'][f'layers.{layer_id}.exit_norm.bias'] = final_norm_bias
            if not use_pipeline_parallel:
                checkpoint_save_path = os.path.join(checkpoint_save_dir, f'mp_rank_{tensor_rank:02d}', 'model_optim_rng.pt')
            else:
                checkpoint_save_path = os.path.join(checkpoint_save_dir, f'mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}', 'model_optim_rng.pt')
            dirname = os.path.dirname(checkpoint_save_path)
            os.makedirs(dirname, exist_ok = True)
            print(f'Saving checkpoint [pp:{pipeline_rank}, tp:{tensor_rank}] to {checkpoint_save_path} ...')
            torch.save(checkpoint_dicts[pipeline_rank], checkpoint_save_path)
    print('Add Exit Layers Completed')

def convert(args):
    if args.megatron_path is not None:
        sys.path.insert(0, args.megatron_path)
    checkpoint_load_dir = os.path.join(args.load_dir, 'iter_{:07d}'.format(args.load_iteration))
    checkpoint_save_dir = os.path.join(args.save_dir, 'iter_{:07d}'.format(args.load_iteration))
    if args.conversion_type == 'exit-position':
        change_exit_position(args, checkpoint_load_dir, checkpoint_save_dir)
    elif args.conversion_type == 'add-exit':
        add_exit(args, checkpoint_load_dir, checkpoint_save_dir)
    with open(os.path.join(args.save_dir, 'latest_checkpointed_iteration.txt'), 'w', encoding='utf-8') as f:
        f.write(str(args.load_iteration))

if __name__ == '__main__':
    args = get_args()
    convert(args)


    elif os.path.exists(os.path.join(checkpoint_root_path, 'mp_rank_00_000')):
        checkpoint_rank_0_dir = 'mp_rank_00_000'
    else:
        raise FileNotFoundError(f'Checkpoint file {checkpoint_root_path} not found')
    checkpoint_path = os.path.join(checkpoint_root_path, checkpoint_rank_0_dir, 'model_optim_rng.pt')
    print(f"Loading args from {checkpoint_root_path}")
    model = torch.load(checkpoint_path)
    return model['args']

# Init method from megatron-lm
def init_method_normal(sigma):

    def init_(tensor):
        return torch.nn.init.normal_(tensor, mean=0.0, std=sigma)

    return init_

def scaled_init_method_normal(sigma, num_layers):
    std = sigma / math.sqrt(2.0 * num_layers)

    def init_(tensor):
        return torch.nn.init.normal_(tensor, mean=0.0, std=std)

    return init_

def change_exit_position(args, checkpoint_load_dir, checkpoint_save_dir):
    checkpoint_args = load_checkpoint_args(checkpoint_load_dir)
    cur_exit_position = 'pre' if checkpoint_args.pre_exit else 'post'
    if cur_exit_position == args.target_exit_position:
        print("No need to convert")
        return
    pipeline_parallel_size = checkpoint_args.pipeline_model_parallel_size
    tensor_parallel_size = checkpoint_args.tensor_model_parallel_size
    exit_layer_nums = checkpoint_args.exit_layer_nums
    if args.target_exit_position == 'pre':
        exit_layer_nums = [layer_num + 1 for layer_num in exit_layer_nums]
    else:
        exit_layer_nums = [layer_num - 1 for layer_num in exit_layer_nums]
    use_pipeline_parallel = pipeline_parallel_size > 1
    for tensor_rank in range(tensor_parallel_size):
        checkpoint_dicts = {}