@ptrblck I benchmarked the performance and got the following results on my dataset:
CPU embeddings, SparseAdam

15 iterations/second

CPU embeddings, Adam

5 seconds/iteration

GPU embeddings, SparseAdam

90 iterations/second

Using CPU embeddings with regular Adam is a 450x slow down. The model becomes untrainable at this point. Using CPU embeddings with SparseAdam is a 6x slow down. While this is not good, it is manageable if it allows us to train a much larger embedding size.
Unfortunately the speed is dependent on the embedding size. If I 4x my embedding size, time/iteration doubles. So I get ~7.5 iterations/second with CPU embeddings + SparseAdam when I 4x embedding size. This limits the usefulness of this solution.
In conclusion, I think there’s a sweet spot where this solution is helpful if you can’t fit your embedding size in GPU memory but the embedding size is not so big that it makes the CPU training unbearable. We’ll keep testing to see if the slowdown is bearable for our particular application.
One interesting to note for anyone who stumbles upon this post, it’s 1.5x faster (in iterations/second) to directly promote the embedding vectors to GPU than wait. E.g. for neural collaborative filtering:
u = u.to(device)
i = i.to(device)
h = torch.cat([u, i], dim = -1)

is faster than:
u = ...
i = ...

h = torch.cat([u, i], dim = -1)
h = h.to(device)