Parameter missing from state_dict

username: Ameet_SD
cooked: <p>TLDR:<br>
When I print out all the IDs present in <code>optimizer_state['param_groups']</code>, it seems like one of them is missing in <code>optimizer_state['state']</code> (index 136).</p>
<p>I am using a modified version of HuggingFace code to load a model from a checkpoint on TPU and run into the following error.</p>
<blockquote>
<p>Traceback (most recent call last):<br>
File “transformers/examples/xla_spawn.py”, line 85, in <br>
main()<br>
File “transformers/examples/xla_spawn.py”, line 81, in main<br>
xmp.spawn(mod._mp_fn, args=(), nprocs=args.num_cores)<br>
File “/anaconda3/envs/torch-xla-1.6/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py”, line 292, in spawn<br>
_start_fn(0, pf_cfg, fn, args)<br>
File “/anaconda3/envs/torch-xla-1.6/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py”, line 229, in _start_fn<br>
fn(gindex, *args)<br>
File “/home/asd/source_code/Multilingual/transformers/examples/language-modeling/run_mlm_synthetic.py”, line 486, in _mp_fn<br>
main()<br>
File “/home/asd/source_code/Multilingual/transformers/examples/language-modeling/run_mlm_synthetic.py”, line 460, in main<br>
trainer.train(model_path=model_path)<br>
File “/home/asd/source_code/Multilingual/transformers/src/transformers/trainer_word_modifications.py”, line 666, in train<br>
self._load_optimizer_and_scheduler(model_path)<br>
File “/home/asd/source_code/Multilingual/transformers/src/transformers/trainer_word_modifications.py”, line 1003, in _load_optimizer_and_scheduler<br>
self.optimizer.load_state_dict(optimizer_state)<br>
File “/anaconda3/envs/torch-xla-1.6/lib/python3.6/site-packages/torch/optim/optimizer.py”, line 123, in load_state_dict<br>
raise ValueError("loaded state dict contains a parameter group "<br>
ValueError: loaded state dict contains a parameter group that doesn’t match the size of optimizer’s group</p>
</blockquote>
<p>After some poking around, I’ve attributed the error to the missing index in <code>optimizer_state['state']</code>. (PFA image). Any reason why this could be happening?</p>
<p>I’m using <code>xm.save(self.optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))</code> to save the model, and <code>self.optimizer.load_state_dict(optimizer_state)</code> to load the model.<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/c/6c5571216dbcd573521522c30c333bac935bdcce.png" data-download-href="https://discuss.pytorch.org/uploads/default/6c5571216dbcd573521522c30c333bac935bdcce" title="Screen Shot 2021-04-19 at 5.20.41 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/c/6c5571216dbcd573521522c30c333bac935bdcce_2_690x185.png" alt="Screen Shot 2021-04-19 at 5.20.41 PM" data-base62-sha1="fsmCLxljEAPDThcKXkhhUQHA8JU" width="690" height="185" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/c/6c5571216dbcd573521522c30c333bac935bdcce_2_690x185.png, https://discuss.pytorch.org/uploads/default/optimized/3X/6/c/6c5571216dbcd573521522c30c333bac935bdcce_2_1035x277.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/c/6c5571216dbcd573521522c30c333bac935bdcce_2_1380x370.png 2x" data-dominant-color="DADAD9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2021-04-19 at 5.20.41 PM</span><span class="informations">2546×684 148 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>