Repository: huggingface/accelerate

Issue:
Error when pickling accelerated optimizers with PyTorch 1.13
### System Info

```Shell
$ pip list
Package                  Version
------------------------ ----------
accelerate               0.13.2
numpy                    1.23.4
nvidia-cublas-cu11       11.10.3.66
nvidia-cuda-nvrtc-cu11   11.7.99
nvidia-cuda-runtime-cu11 11.7.99
nvidia-cudnn-cu11        8.5.0.96
packaging                21.3
pip                      22.3
psutil                   5.9.3
pyparsing                3.0.9
PyYAML                   6.0
setuptools               65.5.0
torch                    1.13.0
typing_extensions        4.4.0
wheel                    0.37.1
```

Same issue with torch-cpu.

OS: Ubuntu 20.04

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)
- [X] My own task or dataset (give details below)

### Reproduction

```python
import pickle
import accelerate
import torch

model = torch.nn.Linear(10, 10)
optimizer = torch.optim.SGD(model.parameters(), 0.1)
pickle.loads(pickle.dumps(optimizer))  # works

accelerator = accelerate.Accelerator(mixed_precision='fp16')
optimizer2 = accelerator.prepare(optimizer)
pickle.loads(pickle.dumps(optimizer2))  # fails
```

### Expected behavior

I ran into a problem with accelerate and PyTorch 1.13. It appears that optimizers cannot be pickled anymore after being accelerated. When running the attached script, I get:

```
Traceback (most recent call last):
  File ".../skorch/foo.py", line 12, in <module>
    pickle.loads(pickle.dumps(optimizer2))
  File ".../torch/optim/optimizer.py", line 84, in __setstate__
    self.defaults.setdefault('differentiable', False)
  File ".../accelerate/optimizer.py", line 90, in defaults
    return self.optimizer.defaults
AttributeError: 'AcceleratedOptimizer' object has no attribute 'optimizer'
```

The offending addition on PyTorch seems to be this line:

https://github.com/pytorch/pytorch/blob/23fe6c8ca15ec2cf6ea74f93aa91cae343ea534f/torch/optim/optimizer.py#L84

which was not present in PyTorch 1.12. At object creation time, PyTorch now tries to access the `defaults` attribute, which in turn calls the `defaults` property in accelerate, which requires the `optimizer` attribute, which doesn't exist and thus errors. At first glance, it looks like `AcceleratedOptimizer` might need its own `__getstate__` and `__setstate__` to solve this but I'm not sure.



Discussion:
Thanks for the great report!
Would you have time to look at this @pacman100 ?

Resolving patch:
diff --git a/src/accelerate/optimizer.py b/src/accelerate/optimizer.py
--- a/src/accelerate/optimizer.py
+++ b/src/accelerate/optimizer.py
@@ -157,3 +157,9 @@ def is_overflow(self):
     def step_was_skipped(self):
         """Whether or not the optimizer step was skipped."""
         return self._is_overflow
+
+    def __getstate__(self):
+        return self.__dict__.copy()
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)


Test patch:
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
new file mode 100644
--- /dev/null
+++ b/tests/test_optimizer.py
@@ -0,0 +1,36 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pickle
+import unittest
+
+import torch
+
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState
+from accelerate.test_utils import require_cpu
+
+
+@require_cpu
+class OptimizerTester(unittest.TestCase):
+    def test_accelerated_optimizer_pickling(self):
+        model = torch.nn.Linear(10, 10)
+        optimizer = torch.optim.SGD(model.parameters(), 0.1)
+        accelerator = Accelerator()
+        optimizer = accelerator.prepare(optimizer)
+        try:
+            pickle.loads(pickle.dumps(optimizer))
+        except Exception as e:
+            self.fail(f"Accelerated optimizer pickling failed with {e}")
+        AcceleratorState._reset_state()