Problem with model accuracy (after restore) on TPU

username: Grankin
cooked: <p>I’m training GPT-2 from huggingface/transformers on TPU. It’s training well. At the end of a training I’ve got loss around 4.36. When I save and restore the model - the loss skyrockets somewhere to 9.75.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/f/afe9a6e9e511e8bedd41a133ecd80a65ac4c233e.png" data-download-href="https://discuss.pytorch.org/uploads/default/afe9a6e9e511e8bedd41a133ecd80a65ac4c233e" title="1"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/f/afe9a6e9e511e8bedd41a133ecd80a65ac4c233e_2_690x282.png" alt="1" data-base62-sha1="p6c9zGfFB03RanSGCsxqltJTstg" width="690" height="282" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/f/afe9a6e9e511e8bedd41a133ecd80a65ac4c233e_2_690x282.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/f/afe9a6e9e511e8bedd41a133ecd80a65ac4c233e_2_1035x423.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/f/afe9a6e9e511e8bedd41a133ecd80a65ac4c233e_2_1380x564.png 2x" data-dominant-color="0B0B0B"><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">1</span><span class="informations">1631×667 59.1 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>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/2/c2a852fcc83b91294cad92a24e167fcca219885e.png" data-download-href="https://discuss.pytorch.org/uploads/default/c2a852fcc83b91294cad92a24e167fcca219885e" title="2"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/2/c2a852fcc83b91294cad92a24e167fcca219885e_2_690x124.png" alt="2" data-base62-sha1="rM1gfEFQQhBUIYFBdom1KnxVj6S" width="690" height="124" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/2/c2a852fcc83b91294cad92a24e167fcca219885e_2_690x124.png, https://discuss.pytorch.org/uploads/default/optimized/3X/c/2/c2a852fcc83b91294cad92a24e167fcca219885e_2_1035x186.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/c/2/c2a852fcc83b91294cad92a24e167fcca219885e_2_1380x248.png 2x" data-dominant-color="0E0E0E"><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">2</span><span class="informations">1722×310 33.5 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>
<p>I’ve got no similar issues with saving and loading on GPU with that code.</p>
<p>The code what is used to save is just this</p>
<p><code>xm.save(model_to_save.state_dict(), output_model_file)</code></p>
<p><code>xm.save</code> is a convinience what moves tensors from TPU to CPU before saving.</p>
<p>The whole code is here <a href="https://github.com/mgrankin/ru_transformers/blob/master/tpu_lm_finetuning.py" class="inline-onebox" rel="noopener nofollow ugc">ru_transformers/tpu_lm_finetuning.py at master · mgrankin/ru_transformers · GitHub</a></p>
<p>I’ve tried to do the following.</p>
<ol>
<li>I’ve tried to do save and load right after the training</li>
</ol>
<pre><code class="lang-auto">    results = evaluate(args, model, tokenizer, "checkpoint-0", False)
    log_info(f"Eval1 {results}")
    
    model = model_class.from_pretrained(args.model_name_or_path, from_tf=bool('.ckpt' in args.model_name_or_path), config=config)
    model.to(args.device)
    results = evaluate(args, model, tokenizer, "checkpoint-0", False)
    log_info(f"Eval2 {results}")
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/5/a57d4e8daba636ba8ac639b4887a92c5d4282ab3.png" data-download-href="https://discuss.pytorch.org/uploads/default/a57d4e8daba636ba8ac639b4887a92c5d4282ab3" title="3"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/5/a57d4e8daba636ba8ac639b4887a92c5d4282ab3_2_690x300.png" alt="3" data-base62-sha1="nBZfA2iFPIVfjL7w6XtxWSlxCr9" width="690" height="300" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/5/a57d4e8daba636ba8ac639b4887a92c5d4282ab3_2_690x300.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/5/a57d4e8daba636ba8ac639b4887a92c5d4282ab3_2_1035x450.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/5/a57d4e8daba636ba8ac639b4887a92c5d4282ab3_2_1380x600.png 2x" data-dominant-color="121212"><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">3</span><span class="informations">1670×728 88.3 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>
<p>Eval2 is much bigger that Eval1</p>
<ol start="2">
<li>I also tried not to recreate the model, but to replace model state_dict with saved state_dict</li>
</ol>
<pre><code class="lang-auto">    results = evaluate(args, model, tokenizer, "checkpoint-0", False)
    log_info(f"Eval1 {results}")
    
    model.load_state_dict(torch.load('output/classic_s/pytorch_model.bin'))
    model.to(args.device)
    results = evaluate(args, model, tokenizer, "checkpoint-0", False)
    log_info(f"Eval2 {results}")
</code></pre>
<p>In that case Eval2 is equal to Eval1</p>
<p>So, there is something that isn’t in a state_dict, but it affects the model perfomance. What can that be?</p>

username: ptrblck
cooked: <p>Skimming through your code, it looks like you are using <code>AdamW</code> as your optimizer, which uses internal states. To be able to properly resume your training, you should also store/restore the optimizer’s <code>state_dict</code>.</p>

username: Grankin
cooked: <p>Thank you for a valuable advice on saving the state of AdamW. But that is not the root problem here. I don’t run train after save/load and it’s performing way worse right after loading.</p>
<p>Davide Libenzi is trying hard to help me with the issue here <a href="https://github.com/pytorch/xla/issues/1245" rel="nofollow noopener">https://github.com/pytorch/xla/issues/1245</a></p>

username: Grankin
cooked: <p>The issue is mostly resolved here<br>
<aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/xla/issues/1245#issuecomment-547777404" target="_blank" rel="nofollow noopener">github.com/pytorch/xla</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/xla/issues/1245#issuecomment-547777404" target="_blank" rel="nofollow noopener">Problem with model accuracy (after restore) on TPU</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2019-10-27" data-time="09:05:48" data-timezone="UTC">09:05AM - 27 Oct 19 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/mgrankin" target="_blank" rel="nofollow noopener">
          <img alt="mgrankin" src="https://avatars0.githubusercontent.com/u/3540879?v=4" class="onebox-avatar-inline" width="20" height="20">
          mgrankin
        </a>
      </div>
    </div>
  </div>
</div>

<div class="github-row">
  <p class="github-content">🐛 Bug
I’m training GPT-2 from huggingface/transformers on TPU. It’s training well. At the end of a training I’ve got loss around...</p>
</div>

<div class="labels">
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>