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.
