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.

You want accelerator.release_memory(), check out the doc :wink:

1 Like