Momentum SGD is very slow with large embedding layer

[TomorrowIsAnOtherDay]

Hi, I have an (256 col* 500000row) embedding layer in my model and I found that Momentum SGD is 10x slower than adam optimizer. When I printed the global variables using API :tf.global_variables(), I found a large Momentum variable created by optimizer, does that means Momentum SGD have not supported sparse update yet?

[allenlavoie]

I'm assuming you mean that AdamOptimizer is slower than MomentumOptimizer; Adam applies momentum to the whole variable, Momentum only applies it to slices affected. Does tf.contrib.opt.LazyAdamOptimizer work for you?

[TomorrowIsAnOtherDay]

@allenlavoie sorry, I mean AdamOptimizer is faster than MomentumOptimizer.In my case,
AdamOptimizer can run 11.58 batch per second, while MomentumOptimizer can only run 5 batch per second.I run TF 1.3.1 in K40 GPU.

[allenlavoie]

Can you post code to reproduce? It sounds like this is not a sparsity issue, since MomentumOptimizer handles sparse updates more quickly than Adam (in theory). Maybe a device placement issue?

[TomorrowIsAnOtherDay]

The code is complex and has a lot preprocessing code(using Dataset API and TF data), I can write a test code for it.Could you give some advice to test whether it is a placement issue? I just changed the code from AdamOptimizer to MomentumOptimizer.The ways how the optimizer placed should be same.

[TomorrowIsAnOtherDay]

And tf.contrib.opt.LazyAdamOptimizer can speed up training speed, It really help.

[allenlavoie]

I'd start with log_device_placement; maybe sparse Momentum updates don't have a GPU kernel (or don't have one for the dtype you're using)?

[TomorrowIsAnOtherDay]

It seems sparse Momentum updates don't have a GPU kernel?

[allenlavoie]

Oh, right, Adam's sparse updates are in terms of other ops rather than being fused. Looks like none of the other optimizers have GPU kernels implemented for the sparse updates (they do have GPU kernels for dense updates).

One approach would be to re-implement MomentumOptimizer's sparse updates in terms of the same ops Adam is using rather than using the custom fused kernel. May be difficult without also using Adam's dense update behavior (which is more correct, but would upset lots of people).

Or it may be possible to implement the sparse fused GPU kernel efficiently using CUB.

Or you could make a simple copy of MomentumOptimizer which always does dense updates, since those will run on the GPU and be faster in your case. This is more of a hack, but certainly contributions welcome for either of the other options.

[TomorrowIsAnOtherDay]

Thanks @allenlavoie ~

[sachinprasadhs]

Hi, Have you checked the new tf.keras.optimizers.experimental.SGD, which can be implemented with the large embeddings. Thanks!