Repository: EleutherAI/gpt-neox

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. 




Discussion:
Time for everyone’s favorite game: running a binary search on the commit history until we find out what went wrong 🙄

@dashstander Have you ruled out this being a weirdness of the new Stability cluster by running it on another GPU source?

@haileyschoelkopf what is the most recent commit that you know you’ve successfully saved a model with?
I've *loaded* a model with the most recent commit on `main` (ebd47f6e9b72cbdba51f5393958d0d3bd95c6e64), would need to dig for most recent commit. (probably have saved a model as of https://github.com/EleutherAI/gpt-neox/commit/46b7d82c34bac79dba9beb02149393a6d59efe2c)

maybe relevant (though I'll make a separate issue for it): 
Checkpoints saved using DeeperSpeed can't be loaded using the `deepspeed_main` branch with new upstream Deepspeed. Not sure if this is solely due to the issue Igor mentioned about not saving `iteration` in Adam optimizer state, or if it's some other format change over Deepspeed versions.

traceback for that:
```
Traceback (most recent call last):
  File "/fsx/hailey/deepspeed-main-neox/gpt-neox/evaluate.py", line 76, in <module>
    main()
  File "/fsx/hailey/deepspeed-main-neox/gpt-neox/evaluate.py", line 35, in main
    model, neox_args = setup_for_inference_or_eval(use_cache=False)
  File "/fsx/hailey/deepspeed-main-neox/gpt-neox/megatron/utils.py", line 440, in setup_for_inference_or_eval
    model, _, _ = setup_model_and_optimizer(
  File "/fsx/hailey/deepspeed-main-neox/gpt-neox/megatron/training.py", line 437, in setup_model_and_optimizer
    neox_args.iteration = load_checkpoint(
  File "/fsx/hailey/deepspeed-main-neox/gpt-neox/megatron/checkpointing.py", line 235, in load_checkpoint
    checkpoint_name, state_dict = model.load_checkpoint(
  File "/fsx/shiv/torchtest/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 2647, in load_checkpoint
    load_path, client_states = self._load_checkpoint(load_dir,
  File "/fsx/shiv/torchtest/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 2713, in _load_checkpoint
    self.load_module_state_dict(state_dict=checkpoint['module'],
  File "/fsx/shiv/torchtest/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 2507, in load_module_state_dict
    self.module.load_state_dict(state_dict, # TODO
  File "/fsx/shiv/torchtest/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1620, in load_state_dict
    raise TypeError("Expected state_dict to be dict-like, got {}.".format(type(state_dict)))
```
I should prob make an issue on Deepspeed about whether there's a way to update checkpoints to newer DeepSpeed versions. This just came up, and also isn't too big a deal bc I can still run on old Deepspeed for models trained pre-upstream deepspeed main branch switch, just inconvenient

Actually yeah, I suspect this is bc of a mismatch between Deepspeed version and which NeoX branch you're on. What's the deepspeed version and what branch is this (main?)
I'm on DeeperSpeed and GPT-NeoX main. I actually think this is an(other) `SlurmRunner` issue. It's failing specifically with the version of the config files that are saved as strings in the arguments. Those specifically were giving me trouble when implementing the SlurmRunner and there's Deep(er)Speed logic to basically clean them up. They get written back to strings to get passed in to `srun` as a command line arg, but it's too much of a coincidence for me not to suspect it.

Maybe make another issue about loading the checkpoints? Or comment on #673 ? That seems like something we should all be aware of
Ok, yeah, this is a `SlurmRunner` thing. To make NeoX and the DeepSpeed SlurmRunner compatible I needed to go in and clean up those configs before passing them to `srun`. But now the when the argument dict is parsed [here](https://github.com/EleutherAI/gpt-neox/blob/ebd47f6e9b72cbdba51f5393958d0d3bd95c6e64/megatron/neox_arguments/arguments.py#L381) all of the configs are so clean they're getting parsed into `dict`s instead of being left as strings.

My proposed solution is to check for this when serializing the configs and either use `json.dumps` or `f.write` as needed. I'm about to get kicked out of the cafe I'm in, but I'll make a PR to that effect tonight.
(Made an issue at #732 for the above error I described)!

Resolving patch:
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 @@ def save_ds_checkpoint(iteration, model, neox_args):
         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):


Test patch: