How to remove a model (unprepare) from the accelerator
I am pretraining a big model which produces embeddings. Every n steps I do a validation run and train a small randomly initialized linear probe. I create a new optimizer and a new model for the linear probe. I call accelerator.prepare for the linear probe model and optimizer.I don’t want to save the linear probe weights or optimizer weights to be saved in the accelerator checkpoint. After finishing a validation run I want to remove the probe and it’s optimizer from the accelerator. I also want to drop them from memory.The problem is that without removing the probe from the accelerator, the checkpoints from save_checkpoint balloon in size, adding a new probe every validation run.
2024-04-29T16:21:02Z
You wantaccelerator.release_memory(), check out the doc