Training speed of NMT models about trax HOT 5 CLOSED

This feels slow indeed. But just to know the baseline: could you provide the speed of a comparable pytorch/TF model on the same hardware with the same batch size?

from trax.

Hello,

On a v100 on similar settings using tensor2tensor, it takes less than 30 seconds per 100 batches or 0.3 seconds per batch. My guess is that, reformer has to recompute activations during backpropagation and this involves an overhead. Another thing is that reformer probably works on TPUs much faster (more compute cores or better addition-subtraction optimizations which reformer relies on). In any case, I would like to know if there is any way to speed things up.

Regards.

from trax.

I think that the batches for Reformer may be much larger - so it may take longer per batch, but the difference in speed per token isn't that big? Have you checked what are the exact batch sizes in both cases, do they match?

from trax.

Hi,

Now that you mention it, I noticed that the batch sizes are 256 SENTENCES per batch instead of the default 2048 TOKENS in transformer big.

Assuming an average token size of 24 thats 6144 tokens per batch on average. So thats almost 3x larger batches. I would expect that a V100 would be able to parallellize on larger batches but I may be wrong. If you think this is what the issue is then I am satisfied. :)

from trax.

I think it is, good to have clarified that :).

from trax.