Hi, actually using sparse gradients for the Embedding did not solve the problem. I still have a very long synchronization time between forward and backward pass.
To recap :

I use a single GPU remotely.
I implemented a light weight two-parts neural network : the first part is just the average of the embeddings (text input), the second part is two linear layers (other input). I end up with two embeddings. I compute the cosine similarity between both embeddings. To have negative samples, I compute embeddings between all samples in the batch. For example if I have a batch size of 512, I have a cosine sim matrix of size 512 x 512. This step takes 0.3ms to be computed.
I use the softmax loss over the cosine similarities (along the first dim)
Synchronization time between all streams of all kernels of my GPU takes around 0.8s to be computed for a batch size of 512 (that’s very long). It is dependent on the size of the batch (with a batch size of 16 it is only 0.02s).
Backward and forward steps are quite fast : 0.008s each

Does someone have an idea about how to speed up the synchro time between forward and backward pass please ? What could be wrong ?
Thank you in advance