def set_up_autotuning(encoded_config, overwrite_values):
        config = json.loads(base64.urlsafe_b64decode(encoded_config).decode("utf-8"))
        overwrite_values = overwrite_values if overwrite_values else {}
        for tuning_param in AUTOTUNING_ARGS:
            # TODO: This is for autotuning specifically, may cause surprises for someone with a weird setup
            if tuning_param in config:
                overwrite_values[tuning_param] = config[tuning_param]
        return overwrite_values

set_up_autotuning(encoded_config='eyJ0cmFpbl9iYXRjaF9zaXplIjogNCwgInRyYWluX21pY3JvX2JhdGNoX3NpemVfcGVyX2dwdSI6IDQsICJvcHRpbWl6ZXIiOiB7InR5cGUiOiAic20zIiwgInBhcmFtcyI6IHt9fSwgImZwMTYiOiB7InR5cGUiOiAiZnAxNiIsICJlbmFibGVkIjogdHJ1ZX0sICJ6ZXJvX29wdGltaXphdGlvbiI6IHsic3RhZ2UiOiAwLCAiYWxsZ2F0aGVyX3BhcnRpdGlvbnMiOiB0cnVlLCAicmVkdWNlX3NjYXR0ZXIiOiB0cnVlLCAiYWxsZ2F0aGVyX2J1Y2tldF9zaXplIjogNTAwMDAwMDAwLCAib3ZlcmxhcF9jb21tIjogZmFsc2UsICJyZWR1Y2VfYnVja2V0X3NpemUiOiA1MDAwMDAwMDAsICJjb250aWd1b3VzX2dyYWRpZW50cyI6IGZhbHNlfSwgIndhbGxfY2xvY2tfYnJlYWtkb3duIjogdHJ1ZSwgImNvbW1zX2xvZ2dlciI6IHsiZW5hYmxlZCI6IHRydWUsICJ2ZXJib3NlIjogdHJ1ZSwgInByb2ZfYWxsIjogdHJ1ZSwgImRlYnVnIjogZmFsc2V9fQ==', overwrite_values={'train_iters': 32})

{'train_batch_size': 4, 'train_micro_batch_size_per_gpu': 4, 'optimizer': {'type': 'sm3', 'params': {}}, 'fp16': {'type': 'fp16', 'enabled': True}, 'zero_optimization': {'stage': 0, 'allgather_partitions': True, 'reduce_scatter': True, 'allgather_bucket_size': 500000000, 'overlap_comm': False, 'reduce_bucket_size': 500000000, 'contiguous_gradients': False}, 'wall_clock_breakdown': True, 'comms_logger': {'enabled': True, 'verbose': True, 'prof_all': True, 'debug': False}}

{'train_iters': 32, 'train_batch_size': 4}

{'train_iters': 32, 'train_batch_size': 4, 'train_micro_batch_size_per_gpu': 4}

{'train_iters': 32, 'train_batch_size': 4, 'train_micro_batch_size_per_gpu': 4, 'zero_optimization': {'stage': 0, 'allgather_partitions': True, 'reduce_scatter': True, 'allgather_bucket_size': 500000000, 'overlap_comm': False, 'reduce_bucket_size': 500000000, 'contiguous_gradients': False}}


def run_train_test(monkeypatch, overwrite_values: dict):
    max_train_iters = 32
    checkpoint_args = {"train_iters": max_train_iters}
    overwrite_values = checkpoint_args
    input_args = ["train.py", "tests/config/test_setup.yml"]
    deepspeed_main_args = simulate_deepy_env(monkeypatch, input_args)

    # Train model, whilst patching collect_loss_for_unit_test to track model loss at each step
    loss_per_iteration = []
    with patch(
        "megatron.training.collect_loss_for_unit_test",
        side_effect=lambda x: loss_per_iteration.append(x),
    ):
        train.main(input_args=deepspeed_main_args, overwrite_values=overwrite_values)
        assert (
            len(loss_per_iteration) == max_train_iters
        ), "patching should have collected loss values from each train step"

        # loss should have decreased by now (otherwise increasing the max_steps parameter could have the testcase pass)
        assert min(loss_per_iteration) < loss_per_iteration[0], (
            "training loss should improve within " + str(max_train_iters) + " steps"
        )

run_train_test(monkeypatch={_setattr=[], _setitem=[], _cwd=None, _savesyspath=None}, overwrite_values={'pos_emb': 'rpe'})

{_setattr=[], _setitem=[(environ({'SHELL': '/bin/bash', 'LSCOLORS': 'Gxfxcxdxbxegedabagacad', 'USER_ZDOTDIR': '/home/XXX', 'COLORTERM': 'truecolor', 'LESS': '-R', 'TERM_PROGRAM_VERSION': '3.2a', 'GVM_VERSION': '1.0.22', 'CONDA_EXE': '/local/rcs/XXX/miniforge3/bin/conda', '_CE_M': '', 'TMUX': '/tmp/tmux-19200/default,59951,3', 'PKG_CONFIG_PATH': '/home/XXX/.gvm/pkgsets/go1.19.1/global/overlay/lib/pkgconfig:/home/XXX/.gvm/pkgsets/go1.19.1/global/overlay/lib/pkgconfig:', '_P9K_TTY': '/dev/pts/7', 'GVM_PATH_BACKUP': '/home/XXX/.gvm/bin:/local/rcs/XXX/miniforge3/envs/mal/bin:/local/rcs/XXX/miniforge3/condabin:/home/XXX/.gdrive-downloader:/local/arise/XXX/miniforge3/bin:/home/XXX/.gvm/pkgsets/go1.19.1/global/bin:/home/XXX/.gvm/gos/go1.19.1/bin:/home/XXX/.gvm/pkgsets/go1.19.1/global/overlay/bin:/home/XXX/.gvm/bin:/home/XXX/.vscode-server/cli/servers/Stable-31c37ee8f63491495ac49e43b8544550fbae4533/server/bin/remote-cli:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/home/XXX/.local/bin:/home/XXX/.local/bin:/home/XXX/.local/bin', 'P9K_TTY': 'old', 'LC_FIG_SET_PARENT': '4c022497-5122-4b80-b325-c89bab32302a', 'PWD': '/local/rcs/XXX/code/pytrace-collector/logs/self_collected/tried/EleutherAI+gpt-neox/EleutherAI+gpt-neox', 'LOGNAME': 'XXX', 'XDG_SESSION_TYPE': 'tty', 'CONDA_PREFIX': '/local/rcs/XXX/miniforge3/envs/EleutherAI+gpt-neox', 'VSCODE_GIT_ASKPASS_NODE': '/home/XXX/.vscode-server/cli/servers/Stable-31c37ee8f63491495ac49e43b8544550fbae4533/server/node', 'MOTD_SHOWN': 'pam', 'VSCODE_INJECTION': '1', 'GVM_OVERLAY_PREFIX': '/home/XXX/.gvm/pkgsets/go1.19.1/global/overlay', 'HOME': '/home/XXX', 'LANG': 'en_US.UTF-8', 'DYLD_LIBRARY_PATH': '/home/XXX/.gvm/pkgsets/go1.19.1/global/overlay/lib:/home/XXX/.gvm/pkgsets/go1.19.1/global/overlay/lib', 'LS_COLORS': 'rs=0:di=01;34:ln=01;36:mh=00:pi=40;33:so=01;35:do=01;35:bd=40;33;01:cd=40;33;01:or=40;31;01:mi=00:su=37;41:sg=30;43:ca=30;41:tw=30;42:ow=34;42:st=37;44:ex=01;32:*.tar=01;31:*.tgz=01;31:*.arc=01;31:*.arj=01;31:*.taz=01;31:*.lha=01;31:*.lz4=01;31:*.lzh=01;31:*.lzma=01;31:*.tlz=01;31:*.txz=01;31:*.tzo=01;31:*.t7z=01;31:*.zip=01;31:*.z=01;31:*.dz=01;31:*.gz=01;31:*.lrz=01;31:*.lz=01;31:*.lzo=01;31:*.xz=01;31:*.zst=01;31:*.tzst=01;31:*.bz2=01;31:*.bz=01;31:*.tbz=01;31:*.tbz2=01;31:*.tz=01;31:*.deb=01;31:*.rpm=01;31:*.jar=01;31:*.war=01;31:*.ear=01;31:*.sar=01;31:*.rar=01;31:*.alz=01;31:*.ace=01;31:*.zoo=01;31:*.cpio=01;31:*.7z=01;31:*.rz=01;31:*.cab=01;31:*.wim=01;31:*.swm=01;31:*.dwm=01;31:*.esd=01;31:*.jpg=01;35:*.jpeg=01;35:*.mjpg=01;35:*.mjpeg=01;35:*.gif=01;35:*.bmp=01;35:*.pbm=01;35:*.pgm=01;35:*.ppm=01;35:*.tga=01;35:*.xbm=01;35:*.xpm=01;35:*.tif=01;35:*.tiff=01;35:*.png=01;35:*.svg=01;35:*.svgz=01;35:*.mng=01;35:*.pcx=01;35:*.mov=01;35:*.mpg=01;35:*.mpeg=01;35:*.m2v=01;35:*.mkv=01;35:*.webm=01;35:*.webp=01;35:*.ogm=01;35:*.mp4=01;35:*.m4v=01;35:*.mp4v=01;35:*.vob=01;35:*.qt=01;35:*.nuv=01;35:*.wmv=01;35:*.asf=01;35:*.rm=01;35:*.rmvb=01;35:*.flc=01;35:*.avi=01;35:*.fli=01;35:*.flv=01;35:*.gl=01;35:*.dl=01;35:*.xcf=01;35:*.xwd=01;35:*.yuv=01;35:*.cgm=01;35:*.emf=01;35:*.ogv=01;35:*.ogx=01;35:*.aac=00;36:*.au=00;36:*.flac=00;36:*.m4a=00;36:*.mid=00;36:*.midi=00;36:*.mka=00;36:*.mp3=00;36:*.mpc=00;36:*.ogg=00;36:*.ra=00;36:*.wav=00;36:*.oga=00;36:*.opus=00;36:*.spx=00;36:*.xspf=00;36:', 'gvm_pkgset_name': 'global', 'SSL_CERT_DIR': '/usr/lib/ssl/certs', 'CONDA_PROMPT_MODIFIER': '(EleutherAI+gpt-neox) ', 'GIT_ASKPASS': '/home/XXX/.vscode-server/cli/servers/Stable-31c37ee8f63491495ac49e43b8544550fbae4533/server/extensions/git/dist/askpass.sh', 'GVM_ROOT': '/home/XXX/.gvm', 'SSH_CONNECTION': '127.0.0.1 55664 127.0.0.1 22', 'GOROOT': '/home/XXX/.gvm/gos/go1.19.1', 'NVM_DIR': '/local/rcs/XXX/.nvm', 'VSCODE_GIT_ASKPASS_EXTRA_ARGS': '', 'XDG_SESSION_CLASS': 'user', 'PYTHONPATH': ':/local/rcs/XXX/code/pytrace-collector:/local/rcs/XXX/code/pytrace-collector:/local/rcs/XXX/code/pytrace-collector:/local/rcs/XXX/code/pytrace-collector/logs/self_...*.gif=01;35:*.bmp=01;35:*.pbm=01;35:*.pgm=01;35:*.ppm=01;35:*.tga=01;35:*.xbm=01;35:*.xpm=01;35:*.tif=01;35:*.tiff=01;35:*.png=01;35:*.svg=01;35:*.svgz=01;35:*.mng=01;35:*.pcx=01;35:*.mov=01;35:*.mpg=01;35:*.mpeg=01;35:*.m2v=01;35:*.mkv=01;35:*.webm=01;35:*.webp=01;35:*.ogm=01;35:*.mp4=01;35:*.m4v=01;35:*.mp4v=01;35:*.vob=01;35:*.qt=01;35:*.nuv=01;35:*.wmv=01;35:*.asf=01;35:*.rm=01;35:*.rmvb=01;35:*.flc=01;35:*.avi=01;35:*.fli=01;35:*.flv=01;35:*.gl=01;35:*.dl=01;35:*.xcf=01;35:*.xwd=01;35:*.yuv=01;35:*.cgm=01;35:*.emf=01;35:*.ogv=01;35:*.ogx=01;35:*.aac=00;36:*.au=00;36:*.flac=00;36:*.m4a=00;36:*.mid=00;36:*.midi=00;36:*.mka=00;36:*.mp3=00;36:*.mpc=00;36:*.ogg=00;36:*.ra=00;36:*.wav=00;36:*.oga=00;36:*.opus=00;36:*.spx=00;36:*.xspf=00;36:', 'gvm_pkgset_name': 'global', 'SSL_CERT_DIR': '/usr/lib/ssl/certs', 'CONDA_PROMPT_MODIFIER': '(EleutherAI+gpt-neox) ', 'GIT_ASKPASS': '/home/XXX/.vscode-server/cli/servers/Stable-31c37ee8f63491495ac49e43b8544550fbae4533/server/extensions/git/dist/askpass.sh', 'GVM_ROOT': '/home/XXX/.gvm', 'SSH_CONNECTION': '127.0.0.1 55664 127.0.0.1 22', 'GOROOT': '/home/XXX/.gvm/gos/go1.19.1', 'NVM_DIR': '/local/rcs/XXX/.nvm', 'VSCODE_GIT_ASKPASS_EXTRA_ARGS': '', 'XDG_SESSION_CLASS': 'user', 'PYTHONPATH': ':/local/rcs/XXX/code/pytrace-collector:/local/rcs/XXX/code/pytrace-collector:/local/rcs/XXX/code/pytrace-collector:/local/rcs/XXX/code/pytrace-collector/logs/self_collected/tried/EleutherAI+gpt-neox/EleutherAI+gpt-neox', 'TERM': 'screen', 'ZSH': '/home/XXX/.oh-my-zsh', '_CE_CONDA': '', 'VSCODE_NONCE': 'd0bc7031-48a3-4719-8bb5-ef236ddd0016', 'ZDOTDIR': '/home/XXX', 'USER': 'XXX', 'TMUX_PANE': '%5', 'VSCODE_GIT_IPC_HANDLE': '/run/user/19200/vscode-git-13d67c6199.sock', 'CONDA_SHLVL': '3', 'SHLVL': '3', 'PAGER': 'less', '_P9K_SSH_TTY': '/dev/pts/7', 'XDG_SESSION_ID': '43', 'CONDA_PYTHON_EXE': '/local/rcs/XXX/miniforge3/bin/python', 'LD_LIBRARY_PATH': '/home/XXX/.gvm/pkgsets/go1.19.1/global/overlay/lib:/home/XXX/.gvm/pkgsets/go1.19.1/global/overlay/lib', 'XDG_RUNTIME_DIR': '/run/user/19200', 'SSL_CERT_FILE': '/usr/lib/ssl/certs/ca-certificates.crt', 'SSH_CLIENT': '127.0.0.1 46946 22', 'CONDA_DEFAULT_ENV': 'EleutherAI+gpt-neox', 'P9K_SSH': '1', 'LC_ALL': 'en_US.UTF-8', 'VSCODE_GIT_ASKPASS_MAIN': '/home/XXX/.vscode-server/cli/servers/Stable-31c37ee8f63491495ac49e43b8544550fbae4533/server/extensions/git/dist/askpass-main.js', 'XDG_DATA_DIRS': '/usr/local/share:/usr/share:/var/lib/snapd/desktop', 'BROWSER': '/home/XXX/.vscode-server/cli/servers/Stable-31c37ee8f63491495ac49e43b8544550fbae4533/server/bin/helpers/browser.sh', 'PATH': '/home/XXX/.gdrive-downloader:/local/arise/XXX/miniforge3/bin:/home/XXX/.gvm/pkgsets/go1.19.1/global/bin:/home/XXX/.gvm/gos/go1.19.1/bin:/home/XXX/.gvm/pkgsets/go1.19.1/global/overlay/bin:/home/XXX/.gvm/bin:/local/rcs/XXX/miniforge3/envs/EleutherAI+gpt-neox/bin:/local/rcs/XXX/miniforge3/condabin:/home/XXX/.gdrive-downloader:/local/arise/XXX/miniforge3/bin:/home/XXX/.vscode-server/cli/servers/Stable-31c37ee8f63491495ac49e43b8544550fbae4533/server/bin/remote-cli:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/home/XXX/.local/bin:/home/XXX/.local/bin', 'DBUS_SESSION_BUS_ADDRESS': 'unix:path=/run/user/19200/bus', 'gvm_go_name': 'go1.19.1', 'CONDA_PREFIX_1': '/local/rcs/XXX/miniforge3', 'CONDA_PREFIX_2': '/local/rcs/XXX/miniforge3/envs/mal', 'OLDPWD': '/local/rcs/XXX/code/pytrace-collector', 'GOPATH': '/home/XXX/.gvm/pkgsets/go1.19.1/global', 'TERM_PROGRAM': 'tmux', 'VSCODE_IPC_HOOK_CLI': '/run/user/19200/vscode-ipc-518d6355-acaf-4714-a359-be3fe9f21e09.sock', '_': '/local/rcs/XXX/miniforge3/envs/EleutherAI+gpt-neox/bin/python', 'TORCH_CUDA_ARCH_LIST': '', 'KMP_DUPLICATE_LIB_OK': 'True', 'KMP_INIT_AT_FORK': 'FALSE', 'CUDA_MODULE_LOADING': 'LAZY', 'PYTEST_CURRENT_TEST': 'tests/model/test_model_train.py::test_model_training_options[pos_emb-rpe] (call)', 'WORLD_SIZE': '1', 'RANK': '0'}), 'RANK', <notset>)], _cwd=None, _savesyspath=None}

{'train_iters': 32, 'train_batch_size': 4, 'train_micro_batch_size_per_gpu': 4, 'zero_optimization': {'stage': 0, 'allgather_partitions': True, 'reduce_scatter': True, 'allgather_bucket_size': 500000000, 'overlap_comm': False, 'reduce_bucket_size': 500000000, 'contiguous_gradients': False}}


def run_neox_args_load_test(yaml_files):
    from megatron.neox_arguments import NeoXArgs

    yaml_list = get_configs_with_path(yaml_files)
    args_loaded = NeoXArgs.from_ymls(yaml_list)
    assert isinstance(args_loaded, NeoXArgs)

    # initialize an empty config dictionary to be filled by yamls
    config = dict()

    # iterate of all to be loaded yaml files
    for conf_file_name in yaml_list:

        # load file
        with open(conf_file_name) as conf_file:
            conf = yaml.load(conf_file, Loader=yaml.FullLoader)

        # check for key duplicates and load values
        for conf_key, conf_value in conf.items():
            if conf_key in config:
                raise ValueError(
                    f"Conf file {conf_file_name} has the following duplicate keys with previously loaded file: {conf_key}"
                )

            conf_key_converted = conf_key.replace(
                "-", "_"
            )  # TODO remove replace and update configuration files?
            config[conf_key_converted] = conf_value

    # validate that neox args has the same value as specified in the config (if specified in the config)
    for k, v in config.items():
        neox_args_value = getattr(args_loaded, k)
        assert v == neox_args_value, (
            "loaded neox args value "
            + str(k)
            + " == "
            + str(neox_args_value)
            + " different from config file "
            + str(v)
        )

run_neox_args_load_test(yaml_files=['125M.yml', 'local_setup.yml', 'cpu_mock_config.yml'])

NeoXArgs(distributed_backend='nccl', local_rank=None, rank=None, lazy_mpu_init=False, short_seq_prob=0.1, eod_mask_loss=False, adlr_autoresume=False, adlr_autoresume_interval=1000, seed=1234, onnx_safe=False, deepscale=False, deepscale_config=None, deepspeed_mpi=False, deepspeed_slurm=False, user_script=None, iteration=None, do_train=None, do_valid=None, do_test=None, save_iters=[10000, 20000, 30000, 40000, 50000, 60000, 70000, 80000, 90000, 100000, 110000, 120000, 130000, 140000, 150000, 160000, 170000, 180000, 190000, 200000, 210000, 220000, 230000, 240000, 250000, 260000, 270000, 280000, 290000, 300000, 310000], global_num_gpus=1, text_gen_type='unconditional', temperature=0.0, top_p=0.0, top_k=0, return_logits=False, maximum_tokens=64, prompt_end='\n', sample_input_file=None, sample_output_file='samples.txt', num_samples=1, recompute=False, eval_results_prefix='', eval_tasks=None, use_wandb=True, wandb_group=None, wandb_team=None, wandb_project='neox', wandb_host='https://api.wandb.ai', wandb_init_all_ranks=False, git_hash='7a8fa2f0', log_dir='logs', tensorboard_dir='tensorboard', log_interval=100, log_grad_pct_zeros=False, log_param_norm=False, log_grad_norm=False, log_optimizer_states=False, log_gradient_noise_scale=False, gradient_noise_scale_n_batches=5, gradient_noise_scale_cpu_offload=False, pipe_parallel_size=1, model_parallel_size=1, pipe_partition_method='type:transformer|mlp', world_size=None, is_pipe_parallel=True, data_path='data/enwik8/enwik8_text_document', use_shared_fs=True, train_data_paths=None, label_data_paths=None, test_data_paths=None, valid_data_paths=None, train_data_weights=None, valid_data_weights=None, test_data_weights=None, weight_by_num_documents=False, weighted_sampler_alpha=1.0, data_impl='mmap', mmap_warmup=False, save='checkpoints', s3_path=None, s3_chunk_size=104857600, config_files={'125M.yml': '# GPT-2 pretraining setup\n{\n   # parallelism settings ( you will want to change these based on your cluster setup, ideally scheduling pipeline stages\n   # across the node boundaries )\n   "pipe_parallel_size": 1,\n   "model_parallel_size": 1,\n\n   # model settings\n   "num_layers": 12,\n   "hidden_size": 768,\n   "num_attention_heads": 12,\n   "seq_length": 2048,\n   "max_position_embeddings": 2048,\n   "norm": "layernorm",\n   "pos_emb": "rotary",\n   "no_weight_tying": true,\n   "gpt_j_residual": false,\n   "output_layer_parallelism": "column",\n\n   # these should provide some speedup but takes a while to build, set to true if desired\n   "scaled_upper_triang_masked_softmax_fusion": false,\n   "bias_gelu_fusion": false,\n   "rope_fusion": false,\n\n   # init methods\n   "init_method": "small_init",\n   "output_layer_init_method": "wang_init",\n\n\n   # optimizer settings\n   "optimizer": {\n     "type": "Adam",\n     "params": {\n       "lr": 0.0006,\n       "betas": [0.9, 0.95],\n       "eps": 1.0e-8,\n     }\n   },\n   "min_lr": 0.00006,\n\n   # for all zero_optimization options, see https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training\n   "zero_optimization": {\n    "stage": 1,\n    "allgather_partitions": True,\n    "allgather_bucket_size": 500000000,\n    "overlap_comm": True,\n    "reduce_scatter": True,\n    "reduce_bucket_size": 500000000,\n    "contiguous_gradients": True,\n  },\n\n   # batch / data settings\n   "train_micro_batch_size_per_gpu": 4,\n   "data_impl": "mmap",\n\n   # activation checkpointing\n   "checkpoint_activations": true,\n   "checkpoint_num_layers": 1,\n   "partition_activations": true,\n   "synchronize_each_layer": true,\n\n   # regularization\n   "gradient_clipping": 1.0,\n   "weight_decay": 0.1,\n   "hidden_dropout": 0.0,\n   "attention_dropout": 0.0,\n\n   # precision settings\n   "fp16": {\n     "enabled": true,\n     "loss_scale": 0,\n     "loss_scale_window": 1000,\n     "hysteresis": 2,\n     "min_loss_scale": 1\n   },\n\n   # misc. training settings\n   "train_iters": 320000,\n   "lr_decay_iters": 320000,\n   "distributed_backend": "nccl",\n   "lr_decay_style": "cosine",\n   "warmup": 0.01,\n   "checkpoint_factor": 10...\n{\n  "global_num_gpus": 1\n}\n'}, load='checkpoints', checkpoint_validation_with_forward_pass=False, checkpoint_scale='linear', checkpoint_factor=10000, extra_save_iters=None, no_save_optim=False, no_save_rng=False, no_load_optim=False, no_load_rng=False, finetune=False, batch_size=4, train_iters=320000, eval_iters=10, keep_last_n_checkpoints=4, eval_interval=1000, split='969, 30, 1', vocab_file='data/gpt2-vocab.json', merge_file='data/gpt2-merges.txt', num_workers=2, exit_interval=None, attention_dropout=0.0, hidden_dropout=0.0, weight_decay=0.1, checkpoint_activations=True, checkpoint_num_layers=1, deepspeed_activation_checkpointing=True, contiguous_checkpointing=False, checkpoint_in_cpu=False, synchronize_each_layer=True, profile_backward=False, partition_activations=True, gas=1, clip_grad=1.0, hysteresis=2, dynamic_loss_scale=True, loss_scale=None, loss_scale_window=1000.0, min_scale=1.0, char_level_ppl=False, use_mup=False, coord_check=False, save_base_shapes=False, base_shapes_file=None, mup_init_scale=1.0, mup_attn_temp=1.0, mup_output_temp=1.0, mup_embedding_mult=1.0, mup_rp_embedding_mult=1.0, mup_width_scale=2, tokenizer_type='GPT2BPETokenizer', padded_vocab_size=None, optimizer_type='Adam', use_bnb_optimizer=False, zero_stage=1, zero_reduce_scatter=True, zero_contiguous_gradients=True, zero_reduce_bucket_size=500000000, zero_allgather_bucket_size=500000000, lr=0.0006, lr_decay_style='cosine', lr_decay_iters=320000, min_lr=6e-05, warmup=0.01, override_lr_scheduler=False, use_checkpoint_lr_scheduler=False, precision='fp16', num_layers=12, hidden_size=768, num_attention_heads=12, seq_length=2048, max_position_embeddings=2048, norm='layernorm', use_qk_layernorm=False, layernorm_epsilon=1e-05, rms_norm_epsilon=1e-08, scalenorm_epsilon=1e-08, pos_emb='rotary', rpe_num_buckets=32, rpe_max_distance=128, opt_pos_emb_offset=0, no_weight_tying=True, attention_config=['global', 'global', 'global', 'global', 'global', 'global', 'global', 'global', 'global', 'global', 'global', 'global'], sparsity_config={}, num_unique_layers=None, param_sharing_style='grouped', make_vocab_size_divisible_by=128, activation='gelu', scaled_upper_triang_masked_softmax_fusion=False, scaled_masked_softmax_fusion=False, bias_gelu_fusion=False, bias_dropout_fusion=False, rope_fusion=False, fp16_lm_cross_entropy=False, init_method_std=0.02, apply_query_key_layer_scaling=False, use_cpu_initialization=False, attention_softmax_in_fp32=False, rotary_pct=1.0, rotary_emb_base=10000, init_method='small_init', output_layer_init_method='wang_init', gmlp_attn_dim=64, gpt_j_residual=False, gpt_j_tied=False, use_bias_in_norms=True, use_bias_in_attn_linear=True, mlp_type='regular', soft_prompt_tuning=None, output_layer_parallelism='column', deepspeed=True, train_batch_size=4, train_micro_batch_size_per_gpu=4, gradient_accumulation_steps=1, optimizer={'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, scheduler=None, fp32_allreduce=False, prescale_gradients=False, gradient_predivide_factor=1.0, sparse_gradients=False, fp16={'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, bf16=None, amp=None, gradient_clipping=1.0, zero_optimization={'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, curriculum_learning=None, curriculum_seqlen=0, steps_per_print=10, wall_clock_breakdown=True, dump_state=False, flops_profiler=None, communication_data_type=None, autotuning=None, activation_checkpointing=None, sparse_attention=None, data_efficiency=None, tensorboard=None, wandb=None, csv_monitor=None, elasticity=None, comms_logger=None, compression_training=None, checkpoint=None, data_types=None, deepspeed_extra_args=None, hostfile='/mock_path', include=None, exclude=None, num_nodes=-1, num_gpus=None, master_port=29500, master_addr=None, launcher='pdsh', force_multi=False, detect_nvlink_pairs=False, autotuning_run=None, no_ssh_check=False, comment=None, account=None)

{'pipe_parallel_size': 1, 'model_parallel_size': 1}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100, 'steps_per_print': 10}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100, 'steps_per_print': 10, 'keep_last_n_checkpoints': 4}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100, 'steps_per_print': 10, 'keep_last_n_checkpoints': 4, 'wall_clock_breakdown': True}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100, 'steps_per_print': 10, 'keep_last_n_checkpoints': 4, 'wall_clock_breakdown': True, 'hostfile': '/mock_path'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100, 'steps_per_print': 10, 'keep_last_n_checkpoints': 4, 'wall_clock_breakdown': True, 'hostfile': '/mock_path', 'data_path': 'data/enwik8/enwik8_text_document'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100, 'steps_per_print': 10, 'keep_last_n_checkpoints': 4, 'wall_clock_breakdown': True, 'hostfile': '/mock_path', 'data_path': 'data/enwik8/enwik8_text_document', 'vocab_file': 'data/gpt2-vocab.json'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100, 'steps_per_print': 10, 'keep_last_n_checkpoints': 4, 'wall_clock_breakdown': True, 'hostfile': '/mock_path', 'data_path': 'data/enwik8/enwik8_text_document', 'vocab_file': 'data/gpt2-vocab.json', 'merge_file': 'data/gpt2-merges.txt'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100, 'steps_per_print': 10, 'keep_last_n_checkpoints': 4, 'wall_clock_breakdown': True, 'hostfile': '/mock_path', 'data_path': 'data/enwik8/enwik8_text_document', 'vocab_file': 'data/gpt2-vocab.json', 'merge_file': 'data/gpt2-merges.txt', 'save': 'checkpoints'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100, 'steps_per_print': 10, 'keep_last_n_checkpoints': 4, 'wall_clock_breakdown': True, 'hostfile': '/mock_path', 'data_path': 'data/enwik8/enwik8_text_document', 'vocab_file': 'data/gpt2-vocab.json', 'merge_file': 'data/gpt2-merges.txt', 'save': 'checkpoints', 'load': 'checkpoints'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100, 'steps_per_print': 10, 'keep_last_n_checkpoints': 4, 'wall_clock_breakdown': True, 'hostfile': '/mock_path', 'data_path': 'data/enwik8/enwik8_text_document', 'vocab_file': 'data/gpt2-vocab.json', 'merge_file': 'data/gpt2-merges.txt', 'save': 'checkpoints', 'load': 'checkpoints', 'checkpoint_validation_with_forward_pass': False}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100, 'steps_per_print': 10, 'keep_last_n_checkpoints': 4, 'wall_clock_breakdown': True, 'hostfile': '/mock_path', 'data_path': 'data/enwik8/enwik8_text_document', 'vocab_file': 'data/gpt2-vocab.json', 'merge_file': 'data/gpt2-merges.txt', 'save': 'checkpoints', 'load': 'checkpoints', 'checkpoint_validation_with_forward_pass': False, 'tensorboard_dir': 'tensorboard'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100, 'steps_per_print': 10, 'keep_last_n_checkpoints': 4, 'wall_clock_breakdown': True, 'hostfile': '/mock_path', 'data_path': 'data/enwik8/enwik8_text_document', 'vocab_file': 'data/gpt2-vocab.json', 'merge_file': 'data/gpt2-merges.txt', 'save': 'checkpoints', 'load': 'checkpoints', 'checkpoint_validation_with_forward_pass': False, 'tensorboard_dir': 'tensorboard', 'log_dir': 'logs'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100, 'steps_per_print': 10, 'keep_last_n_checkpoints': 4, 'wall_clock_breakdown': True, 'hostfile': '/mock_path', 'data_path': 'data/enwik8/enwik8_text_document', 'vocab_file': 'data/gpt2-vocab.json', 'merge_file': 'data/gpt2-merges.txt', 'save': 'checkpoints', 'load': 'checkpoints', 'checkpoint_validation_with_forward_pass': False, 'tensorboard_dir': 'tensorboard', 'log_dir': 'logs', 'use_wandb': True}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100, 'steps_per_print': 10, 'keep_last_n_checkpoints': 4, 'wall_clock_breakdown': True, 'hostfile': '/mock_path', 'data_path': 'data/enwik8/enwik8_text_document', 'vocab_file': 'data/gpt2-vocab.json', 'merge_file': 'data/gpt2-merges.txt', 'save': 'checkpoints', 'load': 'checkpoints', 'checkpoint_validation_with_forward_pass': False, 'tensorboard_dir': 'tensorboard', 'log_dir': 'logs', 'use_wandb': True, 'wandb_host': 'https://api.wandb.ai'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100, 'steps_per_print': 10, 'keep_last_n_checkpoints': 4, 'wall_clock_breakdown': True, 'hostfile': '/mock_path', 'data_path': 'data/enwik8/enwik8_text_document', 'vocab_file': 'data/gpt2-vocab.json', 'merge_file': 'data/gpt2-merges.txt', 'save': 'checkpoints', 'load': 'checkpoints', 'checkpoint_validation_with_forward_pass': False, 'tensorboard_dir': 'tensorboard', 'log_dir': 'logs', 'use_wandb': True, 'wandb_host': 'https://api.wandb.ai', 'wandb_project': 'neox'}

{'pipe_parallel_size': 1, 'model_parallel_size': 1, 'num_layers': 12, 'hidden_size': 768, 'num_attention_heads': 12, 'seq_length': 2048, 'max_position_embeddings': 2048, 'norm': 'layernorm', 'pos_emb': 'rotary', 'no_weight_tying': True, 'gpt_j_residual': False, 'output_layer_parallelism': 'column', 'scaled_upper_triang_masked_softmax_fusion': False, 'bias_gelu_fusion': False, 'rope_fusion': False, 'init_method': 'small_init', 'output_layer_init_method': 'wang_init', 'optimizer': {'type': 'Adam', 'params': {'lr': 0.0006, 'betas': [0.9, 0.95], 'eps': 1e-08}}, 'min_lr': 6e-05, 'zero_optimization': {'stage': 1, 'allgather_partitions': True, 'allgather_bucket_size': 500000000, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 500000000, 'contiguous_gradients': True}, 'train_micro_batch_size_per_gpu': 4, 'data_impl': 'mmap', 'checkpoint_activations': True, 'checkpoint_num_layers': 1, 'partition_activations': True, 'synchronize_each_layer': True, 'gradient_clipping': 1.0, 'weight_decay': 0.1, 'hidden_dropout': 0.0, 'attention_dropout': 0.0, 'fp16': {'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'hysteresis': 2, 'min_loss_scale': 1}, 'train_iters': 320000, 'lr_decay_iters': 320000, 'distributed_backend': 'nccl', 'lr_decay_style': 'cosine', 'warmup': 0.01, 'checkpoint_factor': 10000, 'eval_interval': 1000, 'eval_iters': 10, 'log_interval': 100, 'steps_per_print': 10, 'keep_last_n_checkpoints': 4, 'wall_clock_breakdown': True, 'hostfile': '/mock_path', 'data_path': 'data/enwik8/enwik8_text_document', 'vocab_file': 'data/gpt2-vocab.json', 'merge_file': 'data/gpt2-merges.txt', 'save': 'checkpoints', 'load': 'checkpoints', 'checkpoint_validation_with_forward_pass': False, 'tensorboard_dir': 'tensorboard', 'log_dir': 'logs', 'use_wandb': True, 'wandb_host': 'https://api.wandb.ai', 'wandb_project': 'neox', 'global_num_gpus': 1}

{'data_path': 'data/enwik8/enwik8_text_document', 'vocab_file': 'data/gpt2-vocab.json', 'merge_file': 'data/gpt2-merges.txt', 'save': 'checkpoints', 'load': 'checkpoints', 'checkpoint_validation_with_forward_pass': False, 'tensorboard_dir': 'tensorboard', 'log_dir': 'logs', 'use_wandb': True, 'wandb_host': 'https://api.wandb.ai', 'wandb_project': 'neox'}