We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Checkpointing fails when using the composer command
## Environment 
Running on colo on branch `moin/bert` rebased off of the latest dev: commit hash `be411d3f71feffcbab84da8671dd530fdbbc2f30`.

## To reproduce

Steps to reproduce the behavior:

1. Checkout `moin/bert` branch
2. Run `composer -n 8 examples/run_mosaic_trainer.py -f composer/yamls/models/glue/sst-2_checkpointing.yaml`

## Expected behavior

If I run `python examples/run_mosaic_trainer.py -f composer/yamls/models/glue/sst-2_checkpointing.yaml` instead of `composer -n 8`, then checkpointing works as expected, and we are able to save a checkpoint successfully and continue training. However, when I run it with `composer -n 8`, we see the following crash (exactly when checkpointing happens): 

```
[BATCH][step=599]: { "lr-AdamW/group0": 0.0000, }
Error in subprocess
----------Subprocess STDOUT----------

----------Subprocess STDERR----------

Command '['/usr/bin/python3.8', '-u', 'examples/run_mosaic_trainer.py', '-f', 'composer/yamls/models/glue/sst-2_checkpointing.yaml']' returned non-zero exit status 1.
Killing subprocess 33667 with SIGTERM
Killing subprocess 33662 with SIGTERM
Killing subprocess 33661 with SIGTERM
Killing subprocess 33669 with SIGTERM
Killing subprocess 33659 with SIGTERM
Waiting 30 seconds for processes to terminate...
```

For reference, checkpointing is set to every 600 iterations. There is no printed output.

## Additional context
This is the next step in getting BERT pre-training working (saving the pre-trained model, and then fine-tuning it on GLUE tasks).


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `composer/trainer/checkpoint.py`
Content:
```
1 # Copyright 2021 MosaicML. All Rights Reserved.
2 
3 import logging
4 import os
5 import random
6 import warnings
7 from typing import Any, Dict, Optional
8 
9 import numpy as np
10 import torch
11 import yaml
12 
13 from composer.core import Event, State
14 from composer.core.types import StateDict
15 from composer.trainer.devices.device import Device
16 from composer.utils import ddp, seed_all
17 
18 log = logging.getLogger(__name__)
19 
20 
21 class CheckpointLoader:
22     """Manager for initializing state and restoring RNG state from existing checkpoints.
23 
24     Args:
25         checkpoint_filepath (str): The path to an existing checkpoint file.
26     """
27 
28     def __init__(self, checkpoint_filepath: str):
29         self.state_dict = torch.load(checkpoint_filepath, map_location='cpu')
30 
31     def load_checkpoint(self, state: State):
32         """Initialize state from the loaded checkpoint's data.
33         """
34 
35         state.load_state_dict(self.state_dict["state"])
36         self.checkpoint_rng_state = self._get_checkpoint_rng_state(state, self.state_dict["rng"])
37 
38         if "seed" in self.state_dict:
39             seed_all(self.state_dict["seed"])
40 
41     def restore_checkpoint_rng_state(self, state: State, device: Device):
42         """Restore the state of all RNG objects in this context from the loaded checkpoint's data.
43         """
44 
45         if self.checkpoint_rng_state is None:
46             return
47 
48         assert ddp.get_world_size() == len(
49             self.checkpoint_rng_state['torch']
50         ), f"invariant violation: if the rng state is being restored, then" \
51             "the world size should be the same as in the checkpoint."
52 
53         torch.set_rng_state(self.checkpoint_rng_state['torch'][ddp.get_global_rank()])
54         device.load_state_dict(self.checkpoint_rng_state['device'][ddp.get_global_rank()])
55         random.setstate(self.checkpoint_rng_state['python'][ddp.get_global_rank()])
56         np.random.set_state(self.checkpoint_rng_state['numpy'][ddp.get_global_rank()])
57 
58         self.checkpoint_rng_state = None
59 
60     def _get_checkpoint_rng_state(self, state: State, checkpoint_rng_state: StateDict) -> Optional[StateDict]:
61         original_world_size = len(checkpoint_rng_state["torch"])
62         if original_world_size == ddp.get_world_size():
63             return checkpoint_rng_state
64         else:
65             warnings.warn(f"The checkpoint was created with world_size({original_world_size}), "
66                           f"which differs from the current world_size({ddp.get_world_size()})."
67                           f"RNG state will not be restored.")
68 
69 
70 class Checkpointer:
71     """Manager for saving state to checkpoint files.
72 
73     Args:
74         checkpoint_folder (str): The path to the folder to store checkpoints in.
75         checkpoint_interval (int): The amount of time units to wait between checkpoints.
76         checkpoint_interval_unit (str): The unit (`"ep"` or `"it"`) that
77             `checkpoint_interval` should be measured in.
78     """
79 
80     def __init__(self, checkpoint_folder: str, checkpoint_interval: int, checkpoint_interval_unit: str):
81         if checkpoint_interval_unit.lower() == "ep":
82             self.save_event = Event.EPOCH_END
83         elif checkpoint_interval_unit.lower() == "it":
84             self.save_event = Event.BATCH_END
85         else:
86             raise ValueError(f"Unknown checkpointing interval: {checkpoint_interval_unit}")
87         self.checkpoint_folder = checkpoint_folder
88         self.save_interval = checkpoint_interval
89 
90     def should_checkpoint(self, state: State, event: Event) -> bool:
91         """Given the current state and event, determine whether a checkpoint needs to be created.
92 
93         Args:
94             state (State): The current State of the trainer.
95             event (Event): The current Event being executed.
96         """
97 
98         if event != self.save_event:
99             return False
100         if self.save_event == Event.EPOCH_END:
101             return state.epoch % self.save_interval == 0
102         if self.save_event == Event.BATCH_END:
103             return state.step % self.save_interval == 0
104         return False
105 
106     def save_checkpoint(self, state: State, seed: int, device: Device, config: Optional[Dict[str, Any]] = None) -> None:
107         """Save the current state to a a new checkpoint file.
108 
109         Args:
110             state (State): The current State of the trainer.
111             device (Device): The Device in use by this process.
112             ddp (DDP): The DDP engine in use by this trainer.
113             config (Optional[Dict[str, Any]]): The hparams used to initialize this trainer, if any.
114         """
115 
116         # Store the rank0 seed, if the seed was provided on trainer init
117         # then this is the same seed on all processes
118         # If the seed was not provided, then the rank0 seed will be copied
119         # to all processes on checkpoint resume.
120         # This will be fixed by: https://github.com/mosaicml/composer/issues/12
121         state_dict = {
122             'state': state.state_dict(),  # should be the same across all ranks. per-rank state not stored
123             'rng': self._get_rng_state(device=device),  # stored across all ranks
124             'seed': seed,
125         }
126         if ddp.get_global_rank() != 0:
127             # only rank 0 saves checkpoints
128             # Need the check down here so all the DDP syncs will work for generating the checkpoint
129             return
130 
131         if config:
132             hparams_path = os.path.join(self.checkpoint_folder, "hparams.yaml")
133             os.makedirs(self.checkpoint_folder, mode=0o775, exist_ok=True)
134             config_yaml_str = yaml.dump(config)
135             try:
136                 with open(hparams_path, "x") as f:
137                     # Storing the config (ex. hparams) in a separate file so they can be modified before resuming
138                     f.write(config_yaml_str)
139             except FileExistsError as e:
140                 with open(hparams_path, "r") as f:
141                     # comparing the parsed hparams to ignore whitespace and formatting differences
142                     if yaml.safe_load(config_yaml_str) != yaml.safe_load(f):
143                         raise RuntimeError(f"The hparams in the existing checkpoint folder {self.checkpoint_folder} "
144                                            "differ from those being used in the current training run. "
145                                            "Please specify a new checkpoint folder.") from e
146         if self.save_event == Event.EPOCH_END:
147             filename = f"ep{state.epoch}.pt"
148         elif self.save_event == Event.BATCH_END:
149             filename = f"it{state.step}.pt"
150         else:
151             raise ValueError(f"Invalid checkpoint event: {self.save_event}")
152         save_file = os.path.join(self.checkpoint_folder, filename)
153         with open(save_file, 'xb') as f:
154             torch.save(state_dict, f)
155         log.info(f'Trainer checkpoint saved to {save_file}')
156 
157     def _get_rng_state(self, device: Device) -> StateDict:
158         rng_state = {
159             "python": ddp.all_gather_object(random.getstate()),
160             "numpy": ddp.all_gather_object(np.random.get_state()),
161             "torch": ddp.all_gather_object(torch.random.get_rng_state()),
162             "device": ddp.all_gather_object(device.state_dict()),
163         }
164         # casting the state dict as on non-rank-0, entries will be None-like
165         return rng_state
166
```
--- 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/composer/trainer/checkpoint.py b/composer/trainer/checkpoint.py
--- a/composer/trainer/checkpoint.py
+++ b/composer/trainer/checkpoint.py
@@ -119,7 +119,6 @@
         # to all processes on checkpoint resume.
         # This will be fixed by: https://github.com/mosaicml/composer/issues/12
         state_dict = {
-            'state': state.state_dict(),  # should be the same across all ranks. per-rank state not stored
             'rng': self._get_rng_state(device=device),  # stored across all ranks
             'seed': seed,
         }
@@ -128,6 +127,9 @@
             # Need the check down here so all the DDP syncs will work for generating the checkpoint
             return
 
+        # we add the state only on rank 0 since other processes don't have loggers to serialize
+        state_dict['state'] = state.state_dict()  # should be the same across all ranks. per-rank state not stored
+
         if config:
             hparams_path = os.path.join(self.checkpoint_folder, "hparams.yaml")
             os.makedirs(self.checkpoint_folder, mode=0o775, exist_ok=True)

{"golden_diff": "diff --git a/composer/trainer/checkpoint.py b/composer/trainer/checkpoint.py\n--- a/composer/trainer/checkpoint.py\n+++ b/composer/trainer/checkpoint.py\n@@ -119,7 +119,6 @@\n         # to all processes on checkpoint resume.\n         # This will be fixed by: https://github.com/mosaicml/composer/issues/12\n         state_dict = {\n-            'state': state.state_dict(),  # should be the same across all ranks. per-rank state not stored\n             'rng': self._get_rng_state(device=device),  # stored across all ranks\n             'seed': seed,\n         }\n@@ -128,6 +127,9 @@\n             # Need the check down here so all the DDP syncs will work for generating the checkpoint\n             return\n \n+        # we add the state only on rank 0 since other processes don't have loggers to serialize\n+        state_dict['state'] = state.state_dict()  # should be the same across all ranks. per-rank state not stored\n+\n         if config:\n             hparams_path = os.path.join(self.checkpoint_folder, \"hparams.yaml\")\n             os.makedirs(self.checkpoint_folder, mode=0o775, exist_ok=True)\n", "issue": "Checkpointing fails when using the composer command\n## Environment \r\nRunning on colo on branch `moin/bert` rebased off of the latest dev: commit hash `be411d3f71feffcbab84da8671dd530fdbbc2f30`.\r\n\r\n## To reproduce\r\n\r\nSteps to reproduce the behavior:\r\n\r\n1. Checkout `moin/bert` branch\r\n2. Run `composer -n 8 examples/run_mosaic_trainer.py -f composer/yamls/models/glue/sst-2_checkpointing.yaml`\r\n\r\n## Expected behavior\r\n\r\nIf I run `python examples/run_mosaic_trainer.py -f composer/yamls/models/glue/sst-2_checkpointing.yaml` instead of `composer -n 8`, then checkpointing works as expected, and we are able to save a checkpoint successfully and continue training. However, when I run it with `composer -n 8`, we see the following crash (exactly when checkpointing happens): \r\n\r\n```\r\n[BATCH][step=599]: { \"lr-AdamW/group0\": 0.0000, }\r\nError in subprocess\r\n----------Subprocess STDOUT----------\r\n\r\n----------Subprocess STDERR----------\r\n\r\nCommand '['/usr/bin/python3.8', '-u', 'examples/run_mosaic_trainer.py', '-f', 'composer/yamls/models/glue/sst-2_checkpointing.yaml']' returned non-zero exit status 1.\r\nKilling subprocess 33667 with SIGTERM\r\nKilling subprocess 33662 with SIGTERM\r\nKilling subprocess 33661 with SIGTERM\r\nKilling subprocess 33669 with SIGTERM\r\nKilling subprocess 33659 with SIGTERM\r\nWaiting 30 seconds for processes to terminate...\r\n```\r\n\r\nFor reference, checkpointing is set to every 600 iterations. There is no printed output.\r\n\r\n## Additional context\r\nThis is the next step in getting BERT pre-training working (saving the pre-trained model, and then fine-tuning it on GLUE tasks).\r\n\n", "before_files": [{"content": "# Copyright 2021 MosaicML. All Rights Reserved.\n\nimport logging\nimport os\nimport random\nimport warnings\nfrom typing import Any, Dict, Optional\n\nimport numpy as np\nimport torch\nimport yaml\n\nfrom composer.core import Event, State\nfrom composer.core.types import StateDict\nfrom composer.trainer.devices.device import Device\nfrom composer.utils import ddp, seed_all\n\nlog = logging.getLogger(__name__)\n\n\nclass CheckpointLoader:\n    \"\"\"Manager for initializing state and restoring RNG state from existing checkpoints.\n\n    Args:\n        checkpoint_filepath (str): The path to an existing checkpoint file.\n    \"\"\"\n\n    def __init__(self, checkpoint_filepath: str):\n        self.state_dict = torch.load(checkpoint_filepath, map_location='cpu')\n\n    def load_checkpoint(self, state: State):\n        \"\"\"Initialize state from the loaded checkpoint's data.\n        \"\"\"\n\n        state.load_state_dict(self.state_dict[\"state\"])\n        self.checkpoint_rng_state = self._get_checkpoint_rng_state(state, self.state_dict[\"rng\"])\n\n        if \"seed\" in self.state_dict:\n            seed_all(self.state_dict[\"seed\"])\n\n    def restore_checkpoint_rng_state(self, state: State, device: Device):\n        \"\"\"Restore the state of all RNG objects in this context from the loaded checkpoint's data.\n        \"\"\"\n\n        if self.checkpoint_rng_state is None:\n            return\n\n        assert ddp.get_world_size() == len(\n            self.checkpoint_rng_state['torch']\n        ), f\"invariant violation: if the rng state is being restored, then\" \\\n            \"the world size should be the same as in the checkpoint.\"\n\n        torch.set_rng_state(self.checkpoint_rng_state['torch'][ddp.get_global_rank()])\n        device.load_state_dict(self.checkpoint_rng_state['device'][ddp.get_global_rank()])\n        random.setstate(self.checkpoint_rng_state['python'][ddp.get_global_rank()])\n        np.random.set_state(self.checkpoint_rng_state['numpy'][ddp.get_global_rank()])\n\n        self.checkpoint_rng_state = None\n\n    def _get_checkpoint_rng_state(self, state: State, checkpoint_rng_state: StateDict) -> Optional[StateDict]:\n        original_world_size = len(checkpoint_rng_state[\"torch\"])\n        if original_world_size == ddp.get_world_size():\n            return checkpoint_rng_state\n        else:\n            warnings.warn(f\"The checkpoint was created with world_size({original_world_size}), \"\n                          f\"which differs from the current world_size({ddp.get_world_size()}).\"\n                          f\"RNG state will not be restored.\")\n\n\nclass Checkpointer:\n    \"\"\"Manager for saving state to checkpoint files.\n\n    Args:\n        checkpoint_folder (str): The path to the folder to store checkpoints in.\n        checkpoint_interval (int): The amount of time units to wait between checkpoints.\n        checkpoint_interval_unit (str): The unit (`\"ep\"` or `\"it\"`) that\n            `checkpoint_interval` should be measured in.\n    \"\"\"\n\n    def __init__(self, checkpoint_folder: str, checkpoint_interval: int, checkpoint_interval_unit: str):\n        if checkpoint_interval_unit.lower() == \"ep\":\n            self.save_event = Event.EPOCH_END\n        elif checkpoint_interval_unit.lower() == \"it\":\n            self.save_event = Event.BATCH_END\n        else:\n            raise ValueError(f\"Unknown checkpointing interval: {checkpoint_interval_unit}\")\n        self.checkpoint_folder = checkpoint_folder\n        self.save_interval = checkpoint_interval\n\n    def should_checkpoint(self, state: State, event: Event) -> bool:\n        \"\"\"Given the current state and event, determine whether a checkpoint needs to be created.\n\n        Args:\n            state (State): The current State of the trainer.\n            event (Event): The current Event being executed.\n        \"\"\"\n\n        if event != self.save_event:\n            return False\n        if self.save_event == Event.EPOCH_END:\n            return state.epoch % self.save_interval == 0\n        if self.save_event == Event.BATCH_END:\n            return state.step % self.save_interval == 0\n        return False\n\n    def save_checkpoint(self, state: State, seed: int, device: Device, config: Optional[Dict[str, Any]] = None) -> None:\n        \"\"\"Save the current state to a a new checkpoint file.\n\n        Args:\n            state (State): The current State of the trainer.\n            device (Device): The Device in use by this process.\n            ddp (DDP): The DDP engine in use by this trainer.\n            config (Optional[Dict[str, Any]]): The hparams used to initialize this trainer, if any.\n        \"\"\"\n\n        # Store the rank0 seed, if the seed was provided on trainer init\n        # then this is the same seed on all processes\n        # If the seed was not provided, then the rank0 seed will be copied\n        # to all processes on checkpoint resume.\n        # This will be fixed by: https://github.com/mosaicml/composer/issues/12\n        state_dict = {\n            'state': state.state_dict(),  # should be the same across all ranks. per-rank state not stored\n            'rng': self._get_rng_state(device=device),  # stored across all ranks\n            'seed': seed,\n        }\n        if ddp.get_global_rank() != 0:\n            # only rank 0 saves checkpoints\n            # Need the check down here so all the DDP syncs will work for generating the checkpoint\n            return\n\n        if config:\n            hparams_path = os.path.join(self.checkpoint_folder, \"hparams.yaml\")\n            os.makedirs(self.checkpoint_folder, mode=0o775, exist_ok=True)\n            config_yaml_str = yaml.dump(config)\n            try:\n                with open(hparams_path, \"x\") as f:\n                    # Storing the config (ex. hparams) in a separate file so they can be modified before resuming\n                    f.write(config_yaml_str)\n            except FileExistsError as e:\n                with open(hparams_path, \"r\") as f:\n                    # comparing the parsed hparams to ignore whitespace and formatting differences\n                    if yaml.safe_load(config_yaml_str) != yaml.safe_load(f):\n                        raise RuntimeError(f\"The hparams in the existing checkpoint folder {self.checkpoint_folder} \"\n                                           \"differ from those being used in the current training run. \"\n                                           \"Please specify a new checkpoint folder.\") from e\n        if self.save_event == Event.EPOCH_END:\n            filename = f\"ep{state.epoch}.pt\"\n        elif self.save_event == Event.BATCH_END:\n            filename = f\"it{state.step}.pt\"\n        else:\n            raise ValueError(f\"Invalid checkpoint event: {self.save_event}\")\n        save_file = os.path.join(self.checkpoint_folder, filename)\n        with open(save_file, 'xb') as f:\n            torch.save(state_dict, f)\n        log.info(f'Trainer checkpoint saved to {save_file}')\n\n    def _get_rng_state(self, device: Device) -> StateDict:\n        rng_state = {\n            \"python\": ddp.all_gather_object(random.getstate()),\n            \"numpy\": ddp.all_gather_object(np.random.get_state()),\n            \"torch\": ddp.all_gather_object(torch.random.get_rng_state()),\n            \"device\": ddp.all_gather_object(device.state_dict()),\n        }\n        # casting the state dict as on non-rank-0, entries will be None-like\n        return rng_state\n", "path": "composer/trainer/checkpoint.py"}], "after_files": [{"content": "# Copyright 2021 MosaicML. All Rights Reserved.\n\nimport logging\nimport os\nimport random\nimport warnings\nfrom typing import Any, Dict, Optional\n\nimport numpy as np\nimport torch\nimport yaml\n\nfrom composer.core import Event, State\nfrom composer.core.types import StateDict\nfrom composer.trainer.devices.device import Device\nfrom composer.utils import ddp, seed_all\n\nlog = logging.getLogger(__name__)\n\n\nclass CheckpointLoader:\n    \"\"\"Manager for initializing state and restoring RNG state from existing checkpoints.\n\n    Args:\n        checkpoint_filepath (str): The path to an existing checkpoint file.\n    \"\"\"\n\n    def __init__(self, checkpoint_filepath: str):\n        self.state_dict = torch.load(checkpoint_filepath, map_location='cpu')\n\n    def load_checkpoint(self, state: State):\n        \"\"\"Initialize state from the loaded checkpoint's data.\n        \"\"\"\n\n        state.load_state_dict(self.state_dict[\"state\"])\n        self.checkpoint_rng_state = self._get_checkpoint_rng_state(state, self.state_dict[\"rng\"])\n\n        if \"seed\" in self.state_dict:\n            seed_all(self.state_dict[\"seed\"])\n\n    def restore_checkpoint_rng_state(self, state: State, device: Device):\n        \"\"\"Restore the state of all RNG objects in this context from the loaded checkpoint's data.\n        \"\"\"\n\n        if self.checkpoint_rng_state is None:\n            return\n\n        assert ddp.get_world_size() == len(\n            self.checkpoint_rng_state['torch']\n        ), f\"invariant violation: if the rng state is being restored, then\" \\\n            \"the world size should be the same as in the checkpoint.\"\n\n        torch.set_rng_state(self.checkpoint_rng_state['torch'][ddp.get_global_rank()])\n        device.load_state_dict(self.checkpoint_rng_state['device'][ddp.get_global_rank()])\n        random.setstate(self.checkpoint_rng_state['python'][ddp.get_global_rank()])\n        np.random.set_state(self.checkpoint_rng_state['numpy'][ddp.get_global_rank()])\n\n        self.checkpoint_rng_state = None\n\n    def _get_checkpoint_rng_state(self, state: State, checkpoint_rng_state: StateDict) -> Optional[StateDict]:\n        original_world_size = len(checkpoint_rng_state[\"torch\"])\n        if original_world_size == ddp.get_world_size():\n            return checkpoint_rng_state\n        else:\n            warnings.warn(f\"The checkpoint was created with world_size({original_world_size}), \"\n                          f\"which differs from the current world_size({ddp.get_world_size()}).\"\n                          f\"RNG state will not be restored.\")\n\n\nclass Checkpointer:\n    \"\"\"Manager for saving state to checkpoint files.\n\n    Args:\n        checkpoint_folder (str): The path to the folder to store checkpoints in.\n        checkpoint_interval (int): The amount of time units to wait between checkpoints.\n        checkpoint_interval_unit (str): The unit (`\"ep\"` or `\"it\"`) that\n            `checkpoint_interval` should be measured in.\n    \"\"\"\n\n    def __init__(self, checkpoint_folder: str, checkpoint_interval: int, checkpoint_interval_unit: str):\n        if checkpoint_interval_unit.lower() == \"ep\":\n            self.save_event = Event.EPOCH_END\n        elif checkpoint_interval_unit.lower() == \"it\":\n            self.save_event = Event.BATCH_END\n        else:\n            raise ValueError(f\"Unknown checkpointing interval: {checkpoint_interval_unit}\")\n        self.checkpoint_folder = checkpoint_folder\n        self.save_interval = checkpoint_interval\n\n    def should_checkpoint(self, state: State, event: Event) -> bool:\n        \"\"\"Given the current state and event, determine whether a checkpoint needs to be created.\n\n        Args:\n            state (State): The current State of the trainer.\n            event (Event): The current Event being executed.\n        \"\"\"\n\n        if event != self.save_event:\n            return False\n        if self.save_event == Event.EPOCH_END:\n            return state.epoch % self.save_interval == 0\n        if self.save_event == Event.BATCH_END:\n            return state.step % self.save_interval == 0\n        return False\n\n    def save_checkpoint(self, state: State, seed: int, device: Device, config: Optional[Dict[str, Any]] = None) -> None:\n        \"\"\"Save the current state to a a new checkpoint file.\n\n        Args:\n            state (State): The current State of the trainer.\n            device (Device): The Device in use by this process.\n            ddp (DDP): The DDP engine in use by this trainer.\n            config (Optional[Dict[str, Any]]): The hparams used to initialize this trainer, if any.\n        \"\"\"\n\n        # Store the rank0 seed, if the seed was provided on trainer init\n        # then this is the same seed on all processes\n        # If the seed was not provided, then the rank0 seed will be copied\n        # to all processes on checkpoint resume.\n        # This will be fixed by: https://github.com/mosaicml/composer/issues/12\n        state_dict = {\n            'rng': self._get_rng_state(device=device),  # stored across all ranks\n            'seed': seed,\n        }\n        if ddp.get_global_rank() != 0:\n            # only rank 0 saves checkpoints\n            # Need the check down here so all the DDP syncs will work for generating the checkpoint\n            return\n\n        # we add the state only on rank 0 since other processes don't have loggers to serialize\n        state_dict['state'] = state.state_dict()  # should be the same across all ranks. per-rank state not stored\n\n        if config:\n            hparams_path = os.path.join(self.checkpoint_folder, \"hparams.yaml\")\n            os.makedirs(self.checkpoint_folder, mode=0o775, exist_ok=True)\n            config_yaml_str = yaml.dump(config)\n            try:\n                with open(hparams_path, \"x\") as f:\n                    # Storing the config (ex. hparams) in a separate file so they can be modified before resuming\n                    f.write(config_yaml_str)\n            except FileExistsError as e:\n                with open(hparams_path, \"r\") as f:\n                    # comparing the parsed hparams to ignore whitespace and formatting differences\n                    if yaml.safe_load(config_yaml_str) != yaml.safe_load(f):\n                        raise RuntimeError(f\"The hparams in the existing checkpoint folder {self.checkpoint_folder} \"\n                                           \"differ from those being used in the current training run. \"\n                                           \"Please specify a new checkpoint folder.\") from e\n        if self.save_event == Event.EPOCH_END:\n            filename = f\"ep{state.epoch}.pt\"\n        elif self.save_event == Event.BATCH_END:\n            filename = f\"it{state.step}.pt\"\n        else:\n            raise ValueError(f\"Invalid checkpoint event: {self.save_event}\")\n        save_file = os.path.join(self.checkpoint_folder, filename)\n        with open(save_file, 'xb') as f:\n            torch.save(state_dict, f)\n        log.info(f'Trainer checkpoint saved to {save_file}')\n\n    def _get_rng_state(self, device: Device) -> StateDict:\n        rng_state = {\n            \"python\": ddp.all_gather_object(random.getstate()),\n            \"numpy\": ddp.all_gather_object(np.random.get_state()),\n            \"torch\": ddp.all_gather_object(torch.random.get_rng_state()),\n            \"device\": ddp.all_gather_object(device.state_dict()),\n        }\n        # casting the state dict as on non-rank-0, entries will be None-like\n        return rng_state\n", "path": "composer/trainer/checkpoint.py"}]}