LayerNorm is extremely time consuming when computing gradient

[tobyyouup]

Hi @lukaszkaiser I use run_metadata to generate timeline for the transformer model, and check which block is time consuming. I have found that layer_norm is extremely cost in gradient computation, as showed in the figure:

The SymbolicGradient in red circle is the gradient computation of the layer norm for each encoder and decoder layer. It nearly cost most of the time. And in feed forward stage, it's not that time consuming, but also heavy, showed in this figure:

I have check the layer norm implementation, and found that you define the "layer_norm_compute" and "layer_norm_compute_grad" function by yourself.

So I don't quite understand why the layer norm is so time consuming, do you have any ideas?

[lukaszkaiser]

One thing to try: could you set allow_defun to False (https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/models/common_layers.py#L33)? That'll break layer norm into separate kernels, and it should be possible to see which one is taking the most time. Thanks for pointing it out!

[tobyyouup]

@lukaszkaiser Thanks, I have set allow_defun to False, and there are so many kernels in layer norm as showed in below figure. Each is not costly, but sum all up cost much. I am wondering why layer norm have broke into so many ops and I am trying to figure it out, do you have any ideas?

Another question is there is a time gap between data feedforward and gradient backprop, as showed in the below figure.

The feedforward cost nearly 100ms, and backprop cost nearly 1000ms, and there is 400ms time gap between feedforward and backprop. And during the time gap, the cpu_pool and cuda_host_bfc increase as showed in the bottom of the figure. I am just use a single GPU. So what happened during the time gap, is it need to transmit data between CPU and GPU?

[lukaszkaiser]

Probably, but it's all a bit strange. Could you try using noam_norm (a very simple op) rather than layer_norm? Is it faster?

[tobyyouup]

@lukaszkaiser Yes, When I change layer_norm to noam_norm, the speed change from global_step/sec: 0.845291 to global_step/sec: 0.912363 for transformer_base on a single gpu.

[lukaszkaiser]

Ok, I think it's time we closed this. It'd be great to make layer_norm faster, but it looks more like a task for upstream TensorFlow -- for now, we have a faster noam_norm and it doesn't look like we can do much more. Except if someone sees a better way, in that case please reopen as tell!