We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Checkpointing fails to save config data because it is a `dict` not a `str`
**Describe the bug**
Saving model checkpoints fails with stack trace:

```
Traceback (most recent call last):^M
  File "/mnt/nvme/home/dashiell/gpt-neox/train.py", line 27, in <module>^M
    pretrain(neox_args=neox_args)^M
  File "/mnt/nvme/home/dashiell/gpt-neox/megatron/training.py", line 106, in pretrain^M
    iteration = train(^M
  File "/mnt/nvme/home/dashiell/gpt-neox/megatron/training.py", line 613, in train^M
    save_checkpoint(^M
  File "/mnt/nvme/home/dashiell/gpt-neox/megatron/checkpointing.py", line 208, in save_checkpoint^M
    save_ds_checkpoint(iteration, model, neox_args)^M
  File "/mnt/nvme/home/dashiell/gpt-neox/megatron/checkpointing.py", line 201, in save_ds_checkpoint^M
    f.write(config_data)^M
TypeError: write() argument must be str, not dict^M
```

**To Reproduce**
`python3 deepy.py --conf_dir configs 1-3B.yml local_setup.yml`


**Expected behavior**
The checkpoint should save without failing. 



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `megatron/checkpointing.py`
Content:
```
1 # Copyright (c) 2021, EleutherAI
2 # This file is based on code by the authors denoted below and has been modified from its original version.
3 #
4 # Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
5 #
6 # Licensed under the Apache License, Version 2.0 (the "License");
7 # you may not use this file except in compliance with the License.
8 # You may obtain a copy of the License at
9 #
10 #     http://www.apache.org/licenses/LICENSE-2.0
11 #
12 # Unless required by applicable law or agreed to in writing, software
13 # distributed under the License is distributed on an "AS IS" BASIS,
14 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15 # See the License for the specific language governing permissions and
16 # limitations under the License.
17 
18 """Input/output checkpointing."""
19 
20 import os
21 import re
22 import shutil
23 import random
24 import sys
25 import numpy as np
26 
27 import torch
28 from glob import glob
29 
30 from megatron import mpu
31 from megatron import print_rank_0
32 from megatron.utils import natural_sort
33 from megatron.text_generation_utils import get_batch, forward_model
34 from pathlib import Path
35 from pprint import pformat
36 
37 
38 def check_checkpoint_args(neox_args, checkpoint_args):
39     """Ensure fixed arguments for a model are the same for the input
40     arguments and the one retrieved from checkpoint."""
41 
42     assert isinstance(checkpoint_args, dict), "args stored in checkpoint is a dict"
43     for checkpoint_arg_name, checkpoint_arg_value in checkpoint_args.items():
44         args_value = getattr(neox_args, checkpoint_arg_name)
45         error_message = "{} value from checkpoint ({}) is not equal to the currently set argument value ({}).".format(
46             checkpoint_arg_name, checkpoint_arg_value, args_value
47         )
48         assert checkpoint_arg_value == args_value, error_message
49 
50 
51 def do_forward_pass(neox_args, model, inference=False):
52 
53     # set to eval mode
54     model_was_in_train = model.training
55     model.eval()
56 
57     # get context tokens
58     # always forward full batch size
59     context_tokens_tensor = (
60         torch.arange(neox_args.seq_length + 1)
61         .repeat((neox_args.train_micro_batch_size_per_gpu, 1))
62         .cuda()
63     )
64 
65     # forward
66     if inference:
67         tokens, attention_mask, position_ids = get_batch(
68             neox_args, context_tokens_tensor[:, : neox_args.seq_length]
69         )
70         model_inputs = (
71             tokens,
72             position_ids,
73             attention_mask,
74             torch.Tensor(),
75         )
76         logits, _ = forward_model(neox_args, model, model_inputs)
77     elif neox_args.is_pipe_parallel:
78         data_iterator = iter([{"text": context_tokens_tensor}])
79         _, logits = model.eval_batch(data_iter=data_iterator, return_logits=True)
80     else:
81         tokens, attention_mask, position_ids = get_batch(
82             neox_args, context_tokens_tensor[:, : neox_args.seq_length]
83         )
84         logits = model((tokens, position_ids, attention_mask))
85 
86     # reset to train mode, if model was in training before
87     if model_was_in_train:
88         model.train()
89 
90     if logits is not None:
91         logits = logits.detach().cpu()[
92             0
93         ]  # just return first batch item (they are all equal)
94 
95     return logits
96 
97 
98 def check_forward_pass(neox_args, model, checkpoint_logits, inference):
99     # do forward pass with loaded checkpoint
100     logits = do_forward_pass(neox_args=neox_args, model=model, inference=inference)
101 
102     # check
103     if (
104         logits is not None and checkpoint_logits is not None
105     ):  # this could be the case for non-final pipeline stages
106         if not (logits == checkpoint_logits).all().item():
107             if mpu.get_data_parallel_rank() == 0:
108                 print(
109                     " > WARNING: validate_checkpoint_forward() forward after load of checkpoint does not yield exactly same result"
110                 )
111             assert (
112                 torch.isclose(logits, checkpoint_logits).all().item()
113             ), "validate_checkpoint_forward() forward after load of checkpoint does not yield a close result"
114 
115 
116 def ensure_directory_exists(filename):
117     """Build filename's path if it does not already exists."""
118     dirname = os.path.dirname(filename)
119     if not os.path.exists(dirname):
120         os.makedirs(dirname)
121 
122 
123 def get_checkpoint_name(checkpoints_path, iteration, release=False, mp_rank=None):
124     """A unified checkpoint name."""
125     if release:
126         directory = "release"
127     else:
128         directory = "iter_{:07d}".format(iteration)
129     return os.path.join(
130         checkpoints_path,
131         directory,
132         "mp_rank_{:02d}".format(
133             mpu.get_model_parallel_rank() if mp_rank is None else mp_rank
134         ),
135         "model_optim_rng.pt",
136     )
137 
138 
139 def delete_old_checkpoints(save_dir, n_to_keep):
140     if torch.distributed.get_rank() == 0:
141         ckpt_dir_regex = r"global_step[\d]*"
142         if save_dir.endswith("/"):
143             save_dir = save_dir.strip("/")
144         all_ckpts = natural_sort(
145             [
146                 i
147                 for i in glob(f"{save_dir}/*")
148                 if os.path.isdir(i) and re.search(ckpt_dir_regex, i)
149             ]
150         )
151         n_to_delete = len(all_ckpts) - n_to_keep
152         if n_to_delete > 0:
153             to_delete = all_ckpts[:n_to_delete]
154             print(f"WARNING: Deleting old checkpoints: \n\t{', '.join(to_delete)}")
155             for ckpt in to_delete:
156                 try:
157                     shutil.rmtree(ckpt)
158                 except FileNotFoundError:
159                     pass
160 
161 
162 def save_ds_checkpoint(iteration, model, neox_args):
163     """Save a model checkpoint."""
164     sd = {
165         "iteration": iteration,
166         "args": {
167             "num_layers": neox_args.num_layers,
168             "hidden_size": neox_args.hidden_size,
169             "num_attention_heads": neox_args.num_attention_heads,
170             "max_position_embeddings": neox_args.max_position_embeddings,
171             "make_vocab_size_divisible_by": neox_args.make_vocab_size_divisible_by,
172             "padded_vocab_size": neox_args.padded_vocab_size,
173             "tokenizer_type": neox_args.tokenizer_type,
174             "model_parallel_size": neox_args.model_parallel_size,
175         },
176     }
177     # rng states.
178     if not neox_args.no_save_rng:
179         sd["random_rng_state"] = random.getstate()
180         sd["np_rng_state"] = np.random.get_state()
181         sd["torch_rng_state"] = torch.get_rng_state()
182         sd["cuda_rng_state"] = torch.cuda.get_rng_state()
183         sd["rng_tracker_states"] = mpu.get_cuda_rng_tracker().get_states()
184 
185     if neox_args.checkpoint_validation_with_forward_pass:
186         logits = do_forward_pass(neox_args=neox_args, model=model)
187         sd["checkpoint_validation_logits"] = logits
188 
189     # checkpoint folder name
190     tag = f"global_step{iteration}"
191 
192     # save checkpoint
193     model.save_checkpoint(neox_args.save, tag=tag, client_state=sd)
194 
195     # save config files
196     if torch.distributed.get_rank() == 0 and neox_args.config_files is not None:
197         configs_directory = os.path.join(neox_args.save, tag, "configs")
198         os.makedirs(configs_directory, exist_ok=True)
199         for config_filename, config_data in neox_args.config_files.items():
200             with open(os.path.join(configs_directory, config_filename), "w") as f:
201                 f.write(config_data)
202 
203 
204 def save_checkpoint(neox_args, iteration, model, optimizer, lr_scheduler):
205     """Save a model checkpoint."""
206 
207     if neox_args.deepspeed:
208         save_ds_checkpoint(iteration, model, neox_args)
209     else:
210         raise ValueError("Must be using deepspeed to use neox")
211 
212     # Wait so everyone is done (necessary)
213     torch.distributed.barrier()
214     if neox_args.keep_last_n_checkpoints is not None:
215         delete_old_checkpoints(neox_args.save, neox_args.keep_last_n_checkpoints)
216 
217     # Wait so everyone is done (not necessary)
218     torch.distributed.barrier()
219 
220 
221 def load_checkpoint(
222     neox_args, model, optimizer, lr_scheduler, inference=False, iteration=None
223 ):
224     """Load a model checkpoint and return the iteration."""
225     if neox_args.deepspeed:
226         load_optim_and_scheduler = (
227             not neox_args.no_load_optim
228         )  # TODO: These should be configured by separate args
229         if neox_args.finetune:
230             load_optim_and_scheduler = False
231         if iteration is not None:
232             tag = f"global_step{iteration}"
233         else:
234             tag = None
235         checkpoint_name, state_dict = model.load_checkpoint(
236             neox_args.load,
237             load_optimizer_states=load_optim_and_scheduler,
238             load_lr_scheduler_states=load_optim_and_scheduler,
239             tag=tag,
240         )
241 
242         if checkpoint_name is None:
243             # if an iteration is specified, we want to raise an error here rather than
244             # continuing silently, since we are trying to load a specific checkpoint
245             if iteration is not None:
246                 available_checkpoints = sorted(
247                     [
248                         int(i.name.replace("global_step", ""))
249                         for i in Path(neox_args.load).glob("global_step*")
250                     ]
251                 )
252                 raise ValueError(
253                     f"Unable to load checkpoint for iteration {iteration}. \nAvailable iterations: {pformat(available_checkpoints)}"
254                 )
255             if mpu.get_data_parallel_rank() == 0:
256                 print("Unable to load checkpoint.")
257 
258             return 0  # iteration 0, if not checkpoint loaded
259     else:
260         raise ValueError("Must be using deepspeed to use neox")
261 
262     # Set iteration.
263     if neox_args.finetune:
264         iteration = 0
265     else:
266         iteration = state_dict.get("iteration") or state_dict.get(
267             "total_iters"
268         )  # total_iters backward compatible with older checkpoints
269         if iteration is None:
270             raise ValueError(
271                 f"Unable to load iteration from checkpoint {checkpoint_name} with keys {state_dict.keys()}, exiting"
272             )
273 
274     # Check arguments.
275     if "args" in state_dict:
276         checkpoint_args = state_dict["args"]
277         check_checkpoint_args(neox_args=neox_args, checkpoint_args=checkpoint_args)
278         print_rank_0(
279             " > validated currently set args with arguments in the checkpoint ..."
280         )
281     else:
282         print_rank_0(" > could not find arguments in the checkpoint for validation...")
283 
284     # Check loaded checkpoint with forward pass
285     if neox_args.checkpoint_validation_with_forward_pass:
286         if "checkpoint_validation_logits" in state_dict:
287             check_forward_pass(
288                 neox_args=neox_args,
289                 model=model,
290                 checkpoint_logits=state_dict["checkpoint_validation_logits"],
291                 inference=inference,
292             )
293             print_rank_0(" > validated loaded checkpoint with forward pass ...")
294         else:
295             if mpu.get_data_parallel_rank() == 0:
296                 print(
297                     " > WARNING: checkpoint_validation_with_forward_pass is configured but no checkpoint validation data available in checkpoint {}".format(
298                         checkpoint_name
299                     )
300                 )
301 
302     # rng states.
303     if not neox_args.finetune and not neox_args.no_load_rng:
304         try:
305             random.setstate(state_dict["random_rng_state"])
306             np.random.set_state(state_dict["np_rng_state"])
307             torch.set_rng_state(state_dict["torch_rng_state"])
308             torch.cuda.set_rng_state(state_dict["cuda_rng_state"])
309             mpu.get_cuda_rng_tracker().set_states(state_dict["rng_tracker_states"])
310         except KeyError:
311             print_rank_0(
312                 "Unable to load optimizer from checkpoint {}. "
313                 "Specify --no-load-rng or --finetune to prevent "
314                 "attempting to load the optimizer state, "
315                 "exiting ...".format(checkpoint_name)
316             )
317             sys.exit()
318 
319     torch.distributed.barrier()
320     if mpu.get_data_parallel_rank() == 0:
321         print("  successfully loaded {}".format(checkpoint_name))
322 
323     return iteration
324
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/megatron/checkpointing.py b/megatron/checkpointing.py
--- a/megatron/checkpointing.py
+++ b/megatron/checkpointing.py
@@ -17,6 +17,7 @@
 
 """Input/output checkpointing."""
 
+import json
 import os
 import re
 import shutil
@@ -198,7 +199,10 @@
         os.makedirs(configs_directory, exist_ok=True)
         for config_filename, config_data in neox_args.config_files.items():
             with open(os.path.join(configs_directory, config_filename), "w") as f:
-                f.write(config_data)
+                if isinstance(config_data, str):
+                    f.write(config_data)
+                else:
+                    json.dump(config_data, f)
 
 
 def save_checkpoint(neox_args, iteration, model, optimizer, lr_scheduler):

{"golden_diff": "diff --git a/megatron/checkpointing.py b/megatron/checkpointing.py\n--- a/megatron/checkpointing.py\n+++ b/megatron/checkpointing.py\n@@ -17,6 +17,7 @@\n \n \"\"\"Input/output checkpointing.\"\"\"\n \n+import json\n import os\n import re\n import shutil\n@@ -198,7 +199,10 @@\n         os.makedirs(configs_directory, exist_ok=True)\n         for config_filename, config_data in neox_args.config_files.items():\n             with open(os.path.join(configs_directory, config_filename), \"w\") as f:\n-                f.write(config_data)\n+                if isinstance(config_data, str):\n+                    f.write(config_data)\n+                else:\n+                    json.dump(config_data, f)\n \n \n def save_checkpoint(neox_args, iteration, model, optimizer, lr_scheduler):\n", "issue": "Checkpointing fails to save config data because it is a `dict` not a `str`\n**Describe the bug**\r\nSaving model checkpoints fails with stack trace:\r\n\r\n```\r\nTraceback (most recent call last):^M\r\n  File \"/mnt/nvme/home/dashiell/gpt-neox/train.py\", line 27, in <module>^M\r\n    pretrain(neox_args=neox_args)^M\r\n  File \"/mnt/nvme/home/dashiell/gpt-neox/megatron/training.py\", line 106, in pretrain^M\r\n    iteration = train(^M\r\n  File \"/mnt/nvme/home/dashiell/gpt-neox/megatron/training.py\", line 613, in train^M\r\n    save_checkpoint(^M\r\n  File \"/mnt/nvme/home/dashiell/gpt-neox/megatron/checkpointing.py\", line 208, in save_checkpoint^M\r\n    save_ds_checkpoint(iteration, model, neox_args)^M\r\n  File \"/mnt/nvme/home/dashiell/gpt-neox/megatron/checkpointing.py\", line 201, in save_ds_checkpoint^M\r\n    f.write(config_data)^M\r\nTypeError: write() argument must be str, not dict^M\r\n```\r\n\r\n**To Reproduce**\r\n`python3 deepy.py --conf_dir configs 1-3B.yml local_setup.yml`\r\n\r\n\r\n**Expected behavior**\r\nThe checkpoint should save without failing. \r\n\r\n\n", "before_files": [{"content": "# Copyright (c) 2021, EleutherAI\n# This file is based on code by the authors denoted below and has been modified from its original version.\n#\n# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Input/output checkpointing.\"\"\"\n\nimport os\nimport re\nimport shutil\nimport random\nimport sys\nimport numpy as np\n\nimport torch\nfrom glob import glob\n\nfrom megatron import mpu\nfrom megatron import print_rank_0\nfrom megatron.utils import natural_sort\nfrom megatron.text_generation_utils import get_batch, forward_model\nfrom pathlib import Path\nfrom pprint import pformat\n\n\ndef check_checkpoint_args(neox_args, checkpoint_args):\n    \"\"\"Ensure fixed arguments for a model are the same for the input\n    arguments and the one retrieved from checkpoint.\"\"\"\n\n    assert isinstance(checkpoint_args, dict), \"args stored in checkpoint is a dict\"\n    for checkpoint_arg_name, checkpoint_arg_value in checkpoint_args.items():\n        args_value = getattr(neox_args, checkpoint_arg_name)\n        error_message = \"{} value from checkpoint ({}) is not equal to the currently set argument value ({}).\".format(\n            checkpoint_arg_name, checkpoint_arg_value, args_value\n        )\n        assert checkpoint_arg_value == args_value, error_message\n\n\ndef do_forward_pass(neox_args, model, inference=False):\n\n    # set to eval mode\n    model_was_in_train = model.training\n    model.eval()\n\n    # get context tokens\n    # always forward full batch size\n    context_tokens_tensor = (\n        torch.arange(neox_args.seq_length + 1)\n        .repeat((neox_args.train_micro_batch_size_per_gpu, 1))\n        .cuda()\n    )\n\n    # forward\n    if inference:\n        tokens, attention_mask, position_ids = get_batch(\n            neox_args, context_tokens_tensor[:, : neox_args.seq_length]\n        )\n        model_inputs = (\n            tokens,\n            position_ids,\n            attention_mask,\n            torch.Tensor(),\n        )\n        logits, _ = forward_model(neox_args, model, model_inputs)\n    elif neox_args.is_pipe_parallel:\n        data_iterator = iter([{\"text\": context_tokens_tensor}])\n        _, logits = model.eval_batch(data_iter=data_iterator, return_logits=True)\n    else:\n        tokens, attention_mask, position_ids = get_batch(\n            neox_args, context_tokens_tensor[:, : neox_args.seq_length]\n        )\n        logits = model((tokens, position_ids, attention_mask))\n\n    # reset to train mode, if model was in training before\n    if model_was_in_train:\n        model.train()\n\n    if logits is not None:\n        logits = logits.detach().cpu()[\n            0\n        ]  # just return first batch item (they are all equal)\n\n    return logits\n\n\ndef check_forward_pass(neox_args, model, checkpoint_logits, inference):\n    # do forward pass with loaded checkpoint\n    logits = do_forward_pass(neox_args=neox_args, model=model, inference=inference)\n\n    # check\n    if (\n        logits is not None and checkpoint_logits is not None\n    ):  # this could be the case for non-final pipeline stages\n        if not (logits == checkpoint_logits).all().item():\n            if mpu.get_data_parallel_rank() == 0:\n                print(\n                    \" > WARNING: validate_checkpoint_forward() forward after load of checkpoint does not yield exactly same result\"\n                )\n            assert (\n                torch.isclose(logits, checkpoint_logits).all().item()\n            ), \"validate_checkpoint_forward() forward after load of checkpoint does not yield a close result\"\n\n\ndef ensure_directory_exists(filename):\n    \"\"\"Build filename's path if it does not already exists.\"\"\"\n    dirname = os.path.dirname(filename)\n    if not os.path.exists(dirname):\n        os.makedirs(dirname)\n\n\ndef get_checkpoint_name(checkpoints_path, iteration, release=False, mp_rank=None):\n    \"\"\"A unified checkpoint name.\"\"\"\n    if release:\n        directory = \"release\"\n    else:\n        directory = \"iter_{:07d}\".format(iteration)\n    return os.path.join(\n        checkpoints_path,\n        directory,\n        \"mp_rank_{:02d}\".format(\n            mpu.get_model_parallel_rank() if mp_rank is None else mp_rank\n        ),\n        \"model_optim_rng.pt\",\n    )\n\n\ndef delete_old_checkpoints(save_dir, n_to_keep):\n    if torch.distributed.get_rank() == 0:\n        ckpt_dir_regex = r\"global_step[\\d]*\"\n        if save_dir.endswith(\"/\"):\n            save_dir = save_dir.strip(\"/\")\n        all_ckpts = natural_sort(\n            [\n                i\n                for i in glob(f\"{save_dir}/*\")\n                if os.path.isdir(i) and re.search(ckpt_dir_regex, i)\n            ]\n        )\n        n_to_delete = len(all_ckpts) - n_to_keep\n        if n_to_delete > 0:\n            to_delete = all_ckpts[:n_to_delete]\n            print(f\"WARNING: Deleting old checkpoints: \\n\\t{', '.join(to_delete)}\")\n            for ckpt in to_delete:\n                try:\n                    shutil.rmtree(ckpt)\n                except FileNotFoundError:\n                    pass\n\n\ndef save_ds_checkpoint(iteration, model, neox_args):\n    \"\"\"Save a model checkpoint.\"\"\"\n    sd = {\n        \"iteration\": iteration,\n        \"args\": {\n            \"num_layers\": neox_args.num_layers,\n            \"hidden_size\": neox_args.hidden_size,\n            \"num_attention_heads\": neox_args.num_attention_heads,\n            \"max_position_embeddings\": neox_args.max_position_embeddings,\n            \"make_vocab_size_divisible_by\": neox_args.make_vocab_size_divisible_by,\n            \"padded_vocab_size\": neox_args.padded_vocab_size,\n            \"tokenizer_type\": neox_args.tokenizer_type,\n            \"model_parallel_size\": neox_args.model_parallel_size,\n        },\n    }\n    # rng states.\n    if not neox_args.no_save_rng:\n        sd[\"random_rng_state\"] = random.getstate()\n        sd[\"np_rng_state\"] = np.random.get_state()\n        sd[\"torch_rng_state\"] = torch.get_rng_state()\n        sd[\"cuda_rng_state\"] = torch.cuda.get_rng_state()\n        sd[\"rng_tracker_states\"] = mpu.get_cuda_rng_tracker().get_states()\n\n    if neox_args.checkpoint_validation_with_forward_pass:\n        logits = do_forward_pass(neox_args=neox_args, model=model)\n        sd[\"checkpoint_validation_logits\"] = logits\n\n    # checkpoint folder name\n    tag = f\"global_step{iteration}\"\n\n    # save checkpoint\n    model.save_checkpoint(neox_args.save, tag=tag, client_state=sd)\n\n    # save config files\n    if torch.distributed.get_rank() == 0 and neox_args.config_files is not None:\n        configs_directory = os.path.join(neox_args.save, tag, \"configs\")\n        os.makedirs(configs_directory, exist_ok=True)\n        for config_filename, config_data in neox_args.config_files.items():\n            with open(os.path.join(configs_directory, config_filename), \"w\") as f:\n                f.write(config_data)\n\n\ndef save_checkpoint(neox_args, iteration, model, optimizer, lr_scheduler):\n    \"\"\"Save a model checkpoint.\"\"\"\n\n    if neox_args.deepspeed:\n        save_ds_checkpoint(iteration, model, neox_args)\n    else:\n        raise ValueError(\"Must be using deepspeed to use neox\")\n\n    # Wait so everyone is done (necessary)\n    torch.distributed.barrier()\n    if neox_args.keep_last_n_checkpoints is not None:\n        delete_old_checkpoints(neox_args.save, neox_args.keep_last_n_checkpoints)\n\n    # Wait so everyone is done (not necessary)\n    torch.distributed.barrier()\n\n\ndef load_checkpoint(\n    neox_args, model, optimizer, lr_scheduler, inference=False, iteration=None\n):\n    \"\"\"Load a model checkpoint and return the iteration.\"\"\"\n    if neox_args.deepspeed:\n        load_optim_and_scheduler = (\n            not neox_args.no_load_optim\n        )  # TODO: These should be configured by separate args\n        if neox_args.finetune:\n            load_optim_and_scheduler = False\n        if iteration is not None:\n            tag = f\"global_step{iteration}\"\n        else:\n            tag = None\n        checkpoint_name, state_dict = model.load_checkpoint(\n            neox_args.load,\n            load_optimizer_states=load_optim_and_scheduler,\n            load_lr_scheduler_states=load_optim_and_scheduler,\n            tag=tag,\n        )\n\n        if checkpoint_name is None:\n            # if an iteration is specified, we want to raise an error here rather than\n            # continuing silently, since we are trying to load a specific checkpoint\n            if iteration is not None:\n                available_checkpoints = sorted(\n                    [\n                        int(i.name.replace(\"global_step\", \"\"))\n                        for i in Path(neox_args.load).glob(\"global_step*\")\n                    ]\n                )\n                raise ValueError(\n                    f\"Unable to load checkpoint for iteration {iteration}. \\nAvailable iterations: {pformat(available_checkpoints)}\"\n                )\n            if mpu.get_data_parallel_rank() == 0:\n                print(\"Unable to load checkpoint.\")\n\n            return 0  # iteration 0, if not checkpoint loaded\n    else:\n        raise ValueError(\"Must be using deepspeed to use neox\")\n\n    # Set iteration.\n    if neox_args.finetune:\n        iteration = 0\n    else:\n        iteration = state_dict.get(\"iteration\") or state_dict.get(\n            \"total_iters\"\n        )  # total_iters backward compatible with older checkpoints\n        if iteration is None:\n            raise ValueError(\n                f\"Unable to load iteration from checkpoint {checkpoint_name} with keys {state_dict.keys()}, exiting\"\n            )\n\n    # Check arguments.\n    if \"args\" in state_dict:\n        checkpoint_args = state_dict[\"args\"]\n        check_checkpoint_args(neox_args=neox_args, checkpoint_args=checkpoint_args)\n        print_rank_0(\n            \" > validated currently set args with arguments in the checkpoint ...\"\n        )\n    else:\n        print_rank_0(\" > could not find arguments in the checkpoint for validation...\")\n\n    # Check loaded checkpoint with forward pass\n    if neox_args.checkpoint_validation_with_forward_pass:\n        if \"checkpoint_validation_logits\" in state_dict:\n            check_forward_pass(\n                neox_args=neox_args,\n                model=model,\n                checkpoint_logits=state_dict[\"checkpoint_validation_logits\"],\n                inference=inference,\n            )\n            print_rank_0(\" > validated loaded checkpoint with forward pass ...\")\n        else:\n            if mpu.get_data_parallel_rank() == 0:\n                print(\n                    \" > WARNING: checkpoint_validation_with_forward_pass is configured but no checkpoint validation data available in checkpoint {}\".format(\n                        checkpoint_name\n                    )\n                )\n\n    # rng states.\n    if not neox_args.finetune and not neox_args.no_load_rng:\n        try:\n            random.setstate(state_dict[\"random_rng_state\"])\n            np.random.set_state(state_dict[\"np_rng_state\"])\n            torch.set_rng_state(state_dict[\"torch_rng_state\"])\n            torch.cuda.set_rng_state(state_dict[\"cuda_rng_state\"])\n            mpu.get_cuda_rng_tracker().set_states(state_dict[\"rng_tracker_states\"])\n        except KeyError:\n            print_rank_0(\n                \"Unable to load optimizer from checkpoint {}. \"\n                \"Specify --no-load-rng or --finetune to prevent \"\n                \"attempting to load the optimizer state, \"\n                \"exiting ...\".format(checkpoint_name)\n            )\n            sys.exit()\n\n    torch.distributed.barrier()\n    if mpu.get_data_parallel_rank() == 0:\n        print(\"  successfully loaded {}\".format(checkpoint_name))\n\n    return iteration\n", "path": "megatron/checkpointing.py"}], "after_files": [{"content": "# Copyright (c) 2021, EleutherAI\n# This file is based on code by the authors denoted below and has been modified from its original version.\n#\n# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Input/output checkpointing.\"\"\"\n\nimport json\nimport os\nimport re\nimport shutil\nimport random\nimport sys\nimport numpy as np\n\nimport torch\nfrom glob import glob\n\nfrom megatron import mpu\nfrom megatron import print_rank_0\nfrom megatron.utils import natural_sort\nfrom megatron.text_generation_utils import get_batch, forward_model\nfrom pathlib import Path\nfrom pprint import pformat\n\n\ndef check_checkpoint_args(neox_args, checkpoint_args):\n    \"\"\"Ensure fixed arguments for a model are the same for the input\n    arguments and the one retrieved from checkpoint.\"\"\"\n\n    assert isinstance(checkpoint_args, dict), \"args stored in checkpoint is a dict\"\n    for checkpoint_arg_name, checkpoint_arg_value in checkpoint_args.items():\n        args_value = getattr(neox_args, checkpoint_arg_name)\n        error_message = \"{} value from checkpoint ({}) is not equal to the currently set argument value ({}).\".format(\n            checkpoint_arg_name, checkpoint_arg_value, args_value\n        )\n        assert checkpoint_arg_value == args_value, error_message\n\n\ndef do_forward_pass(neox_args, model, inference=False):\n\n    # set to eval mode\n    model_was_in_train = model.training\n    model.eval()\n\n    # get context tokens\n    # always forward full batch size\n    context_tokens_tensor = (\n        torch.arange(neox_args.seq_length + 1)\n        .repeat((neox_args.train_micro_batch_size_per_gpu, 1))\n        .cuda()\n    )\n\n    # forward\n    if inference:\n        tokens, attention_mask, position_ids = get_batch(\n            neox_args, context_tokens_tensor[:, : neox_args.seq_length]\n        )\n        model_inputs = (\n            tokens,\n            position_ids,\n            attention_mask,\n            torch.Tensor(),\n        )\n        logits, _ = forward_model(neox_args, model, model_inputs)\n    elif neox_args.is_pipe_parallel:\n        data_iterator = iter([{\"text\": context_tokens_tensor}])\n        _, logits = model.eval_batch(data_iter=data_iterator, return_logits=True)\n    else:\n        tokens, attention_mask, position_ids = get_batch(\n            neox_args, context_tokens_tensor[:, : neox_args.seq_length]\n        )\n        logits = model((tokens, position_ids, attention_mask))\n\n    # reset to train mode, if model was in training before\n    if model_was_in_train:\n        model.train()\n\n    if logits is not None:\n        logits = logits.detach().cpu()[\n            0\n        ]  # just return first batch item (they are all equal)\n\n    return logits\n\n\ndef check_forward_pass(neox_args, model, checkpoint_logits, inference):\n    # do forward pass with loaded checkpoint\n    logits = do_forward_pass(neox_args=neox_args, model=model, inference=inference)\n\n    # check\n    if (\n        logits is not None and checkpoint_logits is not None\n    ):  # this could be the case for non-final pipeline stages\n        if not (logits == checkpoint_logits).all().item():\n            if mpu.get_data_parallel_rank() == 0:\n                print(\n                    \" > WARNING: validate_checkpoint_forward() forward after load of checkpoint does not yield exactly same result\"\n                )\n            assert (\n                torch.isclose(logits, checkpoint_logits).all().item()\n            ), \"validate_checkpoint_forward() forward after load of checkpoint does not yield a close result\"\n\n\ndef ensure_directory_exists(filename):\n    \"\"\"Build filename's path if it does not already exists.\"\"\"\n    dirname = os.path.dirname(filename)\n    if not os.path.exists(dirname):\n        os.makedirs(dirname)\n\n\ndef get_checkpoint_name(checkpoints_path, iteration, release=False, mp_rank=None):\n    \"\"\"A unified checkpoint name.\"\"\"\n    if release:\n        directory = \"release\"\n    else:\n        directory = \"iter_{:07d}\".format(iteration)\n    return os.path.join(\n        checkpoints_path,\n        directory,\n        \"mp_rank_{:02d}\".format(\n            mpu.get_model_parallel_rank() if mp_rank is None else mp_rank\n        ),\n        \"model_optim_rng.pt\",\n    )\n\n\ndef delete_old_checkpoints(save_dir, n_to_keep):\n    if torch.distributed.get_rank() == 0:\n        ckpt_dir_regex = r\"global_step[\\d]*\"\n        if save_dir.endswith(\"/\"):\n            save_dir = save_dir.strip(\"/\")\n        all_ckpts = natural_sort(\n            [\n                i\n                for i in glob(f\"{save_dir}/*\")\n                if os.path.isdir(i) and re.search(ckpt_dir_regex, i)\n            ]\n        )\n        n_to_delete = len(all_ckpts) - n_to_keep\n        if n_to_delete > 0:\n            to_delete = all_ckpts[:n_to_delete]\n            print(f\"WARNING: Deleting old checkpoints: \\n\\t{', '.join(to_delete)}\")\n            for ckpt in to_delete:\n                try:\n                    shutil.rmtree(ckpt)\n                except FileNotFoundError:\n                    pass\n\n\ndef save_ds_checkpoint(iteration, model, neox_args):\n    \"\"\"Save a model checkpoint.\"\"\"\n    sd = {\n        \"iteration\": iteration,\n        \"args\": {\n            \"num_layers\": neox_args.num_layers,\n            \"hidden_size\": neox_args.hidden_size,\n            \"num_attention_heads\": neox_args.num_attention_heads,\n            \"max_position_embeddings\": neox_args.max_position_embeddings,\n            \"make_vocab_size_divisible_by\": neox_args.make_vocab_size_divisible_by,\n            \"padded_vocab_size\": neox_args.padded_vocab_size,\n            \"tokenizer_type\": neox_args.tokenizer_type,\n            \"model_parallel_size\": neox_args.model_parallel_size,\n        },\n    }\n    # rng states.\n    if not neox_args.no_save_rng:\n        sd[\"random_rng_state\"] = random.getstate()\n        sd[\"np_rng_state\"] = np.random.get_state()\n        sd[\"torch_rng_state\"] = torch.get_rng_state()\n        sd[\"cuda_rng_state\"] = torch.cuda.get_rng_state()\n        sd[\"rng_tracker_states\"] = mpu.get_cuda_rng_tracker().get_states()\n\n    if neox_args.checkpoint_validation_with_forward_pass:\n        logits = do_forward_pass(neox_args=neox_args, model=model)\n        sd[\"checkpoint_validation_logits\"] = logits\n\n    # checkpoint folder name\n    tag = f\"global_step{iteration}\"\n\n    # save checkpoint\n    model.save_checkpoint(neox_args.save, tag=tag, client_state=sd)\n\n    # save config files\n    if torch.distributed.get_rank() == 0 and neox_args.config_files is not None:\n        configs_directory = os.path.join(neox_args.save, tag, \"configs\")\n        os.makedirs(configs_directory, exist_ok=True)\n        for config_filename, config_data in neox_args.config_files.items():\n            with open(os.path.join(configs_directory, config_filename), \"w\") as f:\n                if isinstance(config_data, str):\n                    f.write(config_data)\n                else:\n                    json.dump(config_data, f)\n\n\ndef save_checkpoint(neox_args, iteration, model, optimizer, lr_scheduler):\n    \"\"\"Save a model checkpoint.\"\"\"\n\n    if neox_args.deepspeed:\n        save_ds_checkpoint(iteration, model, neox_args)\n    else:\n        raise ValueError(\"Must be using deepspeed to use neox\")\n\n    # Wait so everyone is done (necessary)\n    torch.distributed.barrier()\n    if neox_args.keep_last_n_checkpoints is not None:\n        delete_old_checkpoints(neox_args.save, neox_args.keep_last_n_checkpoints)\n\n    # Wait so everyone is done (not necessary)\n    torch.distributed.barrier()\n\n\ndef load_checkpoint(\n    neox_args, model, optimizer, lr_scheduler, inference=False, iteration=None\n):\n    \"\"\"Load a model checkpoint and return the iteration.\"\"\"\n    if neox_args.deepspeed:\n        load_optim_and_scheduler = (\n            not neox_args.no_load_optim\n        )  # TODO: These should be configured by separate args\n        if neox_args.finetune:\n            load_optim_and_scheduler = False\n        if iteration is not None:\n            tag = f\"global_step{iteration}\"\n        else:\n            tag = None\n        checkpoint_name, state_dict = model.load_checkpoint(\n            neox_args.load,\n            load_optimizer_states=load_optim_and_scheduler,\n            load_lr_scheduler_states=load_optim_and_scheduler,\n            tag=tag,\n        )\n\n        if checkpoint_name is None:\n            # if an iteration is specified, we want to raise an error here rather than\n            # continuing silently, since we are trying to load a specific checkpoint\n            if iteration is not None:\n                available_checkpoints = sorted(\n                    [\n                        int(i.name.replace(\"global_step\", \"\"))\n                        for i in Path(neox_args.load).glob(\"global_step*\")\n                    ]\n                )\n                raise ValueError(\n                    f\"Unable to load checkpoint for iteration {iteration}. \\nAvailable iterations: {pformat(available_checkpoints)}\"\n                )\n            if mpu.get_data_parallel_rank() == 0:\n                print(\"Unable to load checkpoint.\")\n\n            return 0  # iteration 0, if not checkpoint loaded\n    else:\n        raise ValueError(\"Must be using deepspeed to use neox\")\n\n    # Set iteration.\n    if neox_args.finetune:\n        iteration = 0\n    else:\n        iteration = state_dict.get(\"iteration\") or state_dict.get(\n            \"total_iters\"\n        )  # total_iters backward compatible with older checkpoints\n        if iteration is None:\n            raise ValueError(\n                f\"Unable to load iteration from checkpoint {checkpoint_name} with keys {state_dict.keys()}, exiting\"\n            )\n\n    # Check arguments.\n    if \"args\" in state_dict:\n        checkpoint_args = state_dict[\"args\"]\n        check_checkpoint_args(neox_args=neox_args, checkpoint_args=checkpoint_args)\n        print_rank_0(\n            \" > validated currently set args with arguments in the checkpoint ...\"\n        )\n    else:\n        print_rank_0(\" > could not find arguments in the checkpoint for validation...\")\n\n    # Check loaded checkpoint with forward pass\n    if neox_args.checkpoint_validation_with_forward_pass:\n        if \"checkpoint_validation_logits\" in state_dict:\n            check_forward_pass(\n                neox_args=neox_args,\n                model=model,\n                checkpoint_logits=state_dict[\"checkpoint_validation_logits\"],\n                inference=inference,\n            )\n            print_rank_0(\" > validated loaded checkpoint with forward pass ...\")\n        else:\n            if mpu.get_data_parallel_rank() == 0:\n                print(\n                    \" > WARNING: checkpoint_validation_with_forward_pass is configured but no checkpoint validation data available in checkpoint {}\".format(\n                        checkpoint_name\n                    )\n                )\n\n    # rng states.\n    if not neox_args.finetune and not neox_args.no_load_rng:\n        try:\n            random.setstate(state_dict[\"random_rng_state\"])\n            np.random.set_state(state_dict[\"np_rng_state\"])\n            torch.set_rng_state(state_dict[\"torch_rng_state\"])\n            torch.cuda.set_rng_state(state_dict[\"cuda_rng_state\"])\n            mpu.get_cuda_rng_tracker().set_states(state_dict[\"rng_tracker_states\"])\n        except KeyError:\n            print_rank_0(\n                \"Unable to load optimizer from checkpoint {}. \"\n                \"Specify --no-load-rng or --finetune to prevent \"\n                \"attempting to load the optimizer state, \"\n                \"exiting ...\".format(checkpoint_name)\n            )\n            sys.exit()\n\n    torch.distributed.barrier()\n    if mpu.get_data_parallel_rank() == 0:\n        print(\"  successfully loaded {}\".format(checkpoint_name))\n\n    return iteration\n", "path": "megatron/checkpointing.py"}]}