Using FAISS along deepspeed for online negative mining

[igrekun]

I want to update my dataset periodically (lets say every 10k steps) with new sampling strategy.

Ideally I would utilize all GPUs to calculate all image embeddings, free GPU memory and let FAISS quickly calculate neighbours for each sample. The neighbour indices get saved to the file and dataloader now constructs batches using updated neighbours.

How would I do that with DeepSpeed? Is it possible to:

1. pause training and use the model to quickly calculate embeddings and update dataloader
2. let faiss use all gpus to get nearest neighbours and resume training (deallocate memory eaten by DeepSpeed but not currently used to be more exact)