# Saving optimizer

**URL:** </static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%3E%5C
**Category:** 🤗Accelerate\
**Created:** [August 1, 2021, 12:20am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884 "2021-08-01T00:20:20Z")\
**Posts on this page:** 20\
**Page:** 1

<div class="post-metadata">

**Author:** ![gabrielwong1991](https://avatars.discourse-cdn.com/v4/letter/g/d6d6ee/32.png) [@gabrielwong1991](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fgabrielwong1991)%5C
**Post date:** [August 1, 2021, 12:20am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F1 "2021-08-01T00:20:20Z")

</div>

From documentation `prepare()` is used to send model, optimizer, data loaders to each TPU core etc…

If I want to save the model I will unwrap the model first by doing `unwrap_model()`.

Is there a function to unwrap the optimizer? I assume `unwrap_model()` is just a function to detach from all cores? So I can use it also like `unwrap_model(optimizer)`

Thanks

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsgugger)%5C
**Post date:** [August 1, 2021, 7:03am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F2 "2021-08-01T07:03:25Z")

</div>

The wrapper of the optimizer in the Accelerate library does not add anything to its state dictionary, so you can just call `optimizer.state_dict` and save it.

---

<div class="post-metadata">

**Author:** ![gabrielwong1991](https://avatars.discourse-cdn.com/v4/letter/g/d6d6ee/32.png) [@gabrielwong1991](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fgabrielwong1991)%5C
**Post date:** [August 1, 2021, 2:14pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F3 "2021-08-01T14:14:24Z")

</div>

Hi sgugger,

Thanks for your reply, I managed to get the model saved and loaded.

On saving optimizer and scheduler, I did something like:

```auto
            accelerator.print("==> Saving optimizer <==")
            accelerator.save(optimizer.state_dict(), "optimizer.pth.tar")
            accelerator.print("==> Saving scheduler <==")
            accelerator.save(scheduler.state_dict(), "scheduler.pth.tar")

```

Then when I load checkpoint for the optimizer and scheduler right before `prepare()` (same place as I load the model checkpoint), for example:

```auto
        optimizer.load_state_dict(torch.load("optimizer.pth.tar"))
        scheduler.load_state_dict(torch.load("scheduler.pth.tar"))

```

I get `SIGKILL` exception from the tpu indicating OOM…

And I thought why not just load on main process? Then use `prepare()` to push to all process?

```auto
      if accelerator.is_main_process:
        optimizer.load_state_dict(torch.load("optimizer.pth.tar"))
        scheduler.load_state_dict(torch.load("scheduler.pth.tar"))

```

It went okay.

But then when saving the next checkpoint for optimizer and scheduler (model saved okay) I get an error:

```auto
Exception in device=TPU:0: 140247764057056
Traceback (most recent call last):
  File "/usr/local/lib/python3.7/dist-packages/torch_xla/distributed/xla_multiprocessing.py", line 329, in _mp_start_fn
    _start_fn(index, pf_cfg, fn, args)
  File "/usr/local/lib/python3.7/dist-packages/torch_xla/distributed/xla_multiprocessing.py", line 323, in _start_fn
    fn(gindex, *args)
  File "/usr/local/lib/python3.7/dist-packages/accelerate/utils.py", line 305, in __call__
    self.launcher(*args)
  File "<ipython-input-19-05ee9dbbadb6>", line 183, in map_fn
    accelerator.save(optimizer.state_dict(), "optimizer.pth.tar")
  File "/usr/local/lib/python3.7/dist-packages/accelerate/optimizer.py", line 91, in state_dict
    return self.optimizer.state_dict()
  File "/usr/local/lib/python3.7/dist-packages/torch/optim/optimizer.py", line 120, in state_dict
    for k, v in self.state.items()}
  File "/usr/local/lib/python3.7/dist-packages/torch/optim/optimizer.py", line 120, in <dictcomp>
    for k, v in self.state.items()}
KeyError: 140247764057056
Exception in device=TPU:7: Cancelled: From /job:tpu_worker/replica:0/task:0:
Cancelled by TearDown.

```

I am so close to get codes ready to train but out of clues on saving checkpoint… any idea how to save and load optimizer and scheduler? Thanks.

Edit:  
Related question?  
[https://github.com/huggingface/transformers/issues/4963](https://github.com/huggingface/transformers/issues/4963)

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsgugger)%5C
**Post date:** [August 2, 2021, 5:43am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F4 "2021-08-02T05:43:49Z")

</div>

Make sure you have a source install of Accelerate as there was a bug in the optimizer state reloading fixed last week. Then can you try to load the state dictionary on the CPU? That should avoid the OOM:

```auto
optimizer.load_state_dict(torch.load("optimizer.pth.tar", map_location="cpu"))

```

You should load the state in all processes as there is nothing that will synchronize them otherwise.

We also have utilities to make checkpointing easier on the roadmap, so it will soon get easier hopefully!

---

<div class="post-metadata">

**Author:** ![gabrielwong1991](https://avatars.discourse-cdn.com/v4/letter/g/d6d6ee/32.png) [@gabrielwong1991](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fgabrielwong1991)%5C
**Post date:** [August 3, 2021, 11:14pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F5 "2021-08-03T23:14:47Z")

</div>

So I source install and compile Accelerate and XLA:

```auto
! pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.9-cp37-cp37m-linux_x86_64.whl

! pip install git+https://github.com/huggingface/accelerate

```

used `map_location="cpu"` still returns `SIGKILL` OOM exception… 😕

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsgugger)%5C
**Post date:** [August 4, 2021, 6:19am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F6 "2021-08-04T06:19:04Z")

</div>

That’s what we use in the `Trainer`, so I’m a bit out of options. Do you have a minimal reproducer we could investigate?

---

<div class="post-metadata">

**Author:** ![gabrielwong1991](https://avatars.discourse-cdn.com/v4/letter/g/d6d6ee/32.png) [@gabrielwong1991](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fgabrielwong1991)%5C
**Post date:** [August 4, 2021, 5:53pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F7 "2021-08-04T17:53:33Z")

</div>

Hi sugger here is an example colab workbook. Please use a VM with larger ram like 32gb:

`https://colab.research.google.com/drive/1cOiQjaElTXsdxNVDhSMfvauibyNzkImW?usp=sharing`

So if you run the code from there it shall save all weights in the first loop. Then stop the loop and change the flag `FROM_CHECKPOINT` to `True`. rerun the `accelerate_launcher()` it shall then load the weights.

Thanks

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsgugger)%5C
**Post date:** [August 9, 2021, 10:26am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F8 "2021-08-09T10:26:38Z")

</div>

I get the SIGKILL signal at the first loop when executing your notebook.

---

<div class="post-metadata">

**Author:** ![gabrielwong1991](https://avatars.discourse-cdn.com/v4/letter/g/d6d6ee/32.png) [@gabrielwong1991](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fgabrielwong1991)%5C
**Post date:** [August 9, 2021, 12:03pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F9 "2021-08-09T12:03:29Z")

</div>

Hi sgugger, sorry for the inconvenience.

Can you delete all the code in this notebook and replace the code with the `SIGKILL` one?

```auto
https://colab.research.google.com/github/pytorch/xla/blob/master/contrib/colab/mnist-training.ipynb

```

This notebook is from pytorchXLA demo and have an instance with more ram 35gb so it should work (for debugging purpose). The standard colab one only have 15gb.

Thanks

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsgugger)%5C
**Post date:** [August 9, 2021, 12:42pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F10 "2021-08-09T12:42:03Z")

</div>

The RAM is not dependent on the notebook but on the user. I have no access to a notebook with 35Gb of RAM. In general it looks like you are using a model that just uses too much RAM for Google Colab.

---

<div class="post-metadata">

**Author:** ![gabrielwong1991](https://avatars.discourse-cdn.com/v4/letter/g/d6d6ee/32.png) [@gabrielwong1991](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fgabrielwong1991)%5C
**Post date:** [August 9, 2021, 12:48pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F11 "2021-08-09T12:48:55Z")

</div>

Hi sgugger,

I learn this from reading their github, they instantiate a VM with 35gb for demo so their demo won’t crash. `https://github.com/pytorch/xla/issues/1870#issuecomment-610643406`. Since the backend of yours is also pytorchXLA and are debugging a problem so I believe it is ok to use

Thanks

---

<div class="post-metadata">

**Author:** ![gabrielwong1991](https://avatars.discourse-cdn.com/v4/letter/g/d6d6ee/32.png) [@gabrielwong1991](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fgabrielwong1991)%5C
**Post date:** [August 9, 2021, 12:52pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F12 "2021-08-09T12:52:41Z")

</div>

But anyway, if I use this method, I can load and save model, optimizer, scheduler with no problem. But when restarting from checkpoint model loads ok but optimizer, scheduler always shows OOM.

I can try with a larger RAM instance in GCP

---

<div class="post-metadata">

**Author:** ![gabrielwong1991](https://avatars.discourse-cdn.com/v4/letter/g/d6d6ee/32.png) [@gabrielwong1991](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fgabrielwong1991)%5C
**Post date:** [August 10, 2021, 6:44pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F13 "2021-08-10T18:44:35Z")

</div>

Hi sgugger, I have now setup the environment for GCP TPU…

I wonder why in `accelerate config` it does not ask to specify the IP address of the TPU nodes?  
I then run `accelerate test`, not surprisingly I get TPU connection timeout `Failed to connect to client mesh master` because it does not know where the TPU node is.

Background info:  
I setup the VM with pytorch xla image  
`https://github.com/pytorch/xla#-consume-prebuilt-compute-vm-images`  
and added the TPU address

```auto
export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"

```

and ran PytorchXLA test script in their tutorial with TPU acceleration successfully.

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsgugger)%5C
**Post date:** [August 11, 2021, 6:47am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F15 "2021-08-11T06:47:20Z")

</div>

You have to do the export, that does set up the IP address of the TPU. On my side could run successfully my scripts after doing that step.

---

<div class="post-metadata">

**Author:** ![gabrielwong1991](https://avatars.discourse-cdn.com/v4/letter/g/d6d6ee/32.png) [@gabrielwong1991](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fgabrielwong1991)%5C
**Post date:** [August 11, 2021, 5:00pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F16 "2021-08-11T17:00:24Z")

</div>

Thanks it works!

It trains but same problem with loading optimizer when in checkpoint restart. The optimizer file is only 1.1Gb big and when it gets loaded RAM usage shoots up to 64gb! (The VM has 64gb) and sigh… SIGKILL

There must be memory issue somewhere but anyway I will try a larger RAM instance now maybe 128gb.

Would it be where optimizer model is loaded into TPUs memory but when the new weights are loaded the old weights in the model don’t get freed?

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsgugger)%5C
**Post date:** [August 11, 2021, 5:35pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F17 "2021-08-11T17:35:34Z")

</div>

That’s definitely weird. The old weights should be released when the new ones are loaded inside. What’s happening is probably the combination of both weights being temporary in RAM in all the 8 processes at the same time. It’s annoying we have to pass through the CPU for TPUs checkpointing (on GPUs we could jsut load directly on the GPU and avoid this issue).

---

<div class="post-metadata">

**Author:** ![gabrielwong1991](https://avatars.discourse-cdn.com/v4/letter/g/d6d6ee/32.png) [@gabrielwong1991](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fgabrielwong1991)%5C
**Post date:** [August 13, 2021, 12:35am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F18 "2021-08-13T00:35:04Z")

</div>

Thanks very much again sgugger, so here is my update:

1. Before I have the model (`AutoModel.from_pretrained()`) and tokenizer (`AutoTokenizer.from_pretrained()`) saved into the local disk and load from it. It saves the state of the model, optimizer, scheduler. it just won’t reload the states from checkpoint, OOM SIGKILL etc…

2. Now if I download the model, tokenizer from the internet every time, it reloads from checkpoints but when it comes to saving the states it said to have exceptions in all TPUs… I have no idea why it could be the optimizer in all TPU aren’t sync.

This gets me wondered if my code have bug or is just not working, i.e. defining the custom model class like one of the post here before where he failed to reload the states too? I remove the `accelerate` package and run only in CPU, it trains, load and save so this confirms the code is fine…

I am out of clues really. The only thing I could try now is not to use accelerate and rolls back to my prior code with Pytorch XLA. During that time I haven’t write any checkpoint yet so it should be interesting to see what happens. Thanks.

---

<div class="post-metadata">

**Author:** ![gabrielwong1991](https://avatars.discourse-cdn.com/v4/letter/g/d6d6ee/32.png) [@gabrielwong1991](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fgabrielwong1991)%5C
**Post date:** [August 13, 2021, 5:16pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F19 "2021-08-13T17:16:26Z")

</div>

Managed to load and save optimizer and scheduler state with my original Pytorch XLA script!

This confirms some issue with `accelerate`

Let me know if you want to follow up sgugger i can send you both of my XLA and accelerate script! Thanks for your help! 😃

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsgugger)%5C
**Post date:** [August 30, 2021, 6:03pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F20 "2021-08-30T18:03:05Z")

</div>

Thanks for the info! If you can share gists with each version, I’d happily look and try to find the issue when I have some bandwidth!

---

<div class="post-metadata">

**Author:** ![raphaelmerx](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/raphaelmerx/32/56182_2.png) [@raphaelmerx](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fraphaelmerx)%5C
**Post date:** [May 18, 2023, 2:26am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsaving-optimizer%2F8884%2F21 "2023-05-18T02:26:26Z")

</div>

Fyi you can also save optimizer + other objects needed to continue training using `Accelerator.save_state()`: [Checkpointing](https://huggingface.co/docs/accelerate/usage_guides/checkpoint)