Sparse grad sampler for embeddings

[ffuuugor]

Embedding updates are typically very sparse, however our naive implementation treats them as dense and spends a lot of time and memory storing and multiplying zeros. Additionally, I'm not sure if we support nn.Embedding(sparse=True) modules.

We can rewrite grad sampler with sparse computations instead

[elichienxD]

Hi @ffuuugor ,

I'm currently also using torch sparse tensor (coo layout) with opacus DPSGD in our project. We find that opacus is presently not compatible with sparse tensor as the MLP input. I just found this enhancement issue to be relevent and I wonder if there is any update or follow-up?

(below is the code for grad_sample_module.py)

I have tried to change L62 from = grad_sample to += grad_sample in order to avoid errors. However, the model cannot be trained well compared to the dense input. I wonder if there is any suggestion on where else should I modify the code?

Thanks,
Eli