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


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…


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.