New FusedAdam (since #424 commit 8 Aug 2019) Issues

Team,

This huge PR #424 was not squashed and commits in a long period of time trigger lots of side effects.

As an example, if I use FusedAdam as of this commit 4a8c4ac then the memory footprint in GPUs enables to train with a certain batch size (eg 3072 tokens per batch).

On master, I have to reduce the batch size not to get a CUDA OOM.
(eg tokens 2944 per batch use case is OpenNMT-py)

More concerning, during training, using the same config (batch size 2944 tokens) and all other params being equal:

MASTER of Sept 5 2019
[2019-09-05 12:29:50,751 INFO] Step 100/50000; acc: 3.55; ppl: 6812.21; xent: 8.83; lr: 0.00002; 23557/28574 tok/s; 148 sec
[2019-09-05 12:31:02,317 INFO] Step 200/50000; acc: 6.69; ppl: 2112.41; xent: 7.66; lr: 0.00003; 48483/58933 tok/s; 220 sec
[2019-09-05 12:32:14,252 INFO] Step 300/50000; acc: 10.53; ppl: 597.04; xent: 6.39; lr: 0.00005; 48349/58938 tok/s; 292 sec
[2019-09-05 12:33:26,245 INFO] Step 400/50000; acc: 13.50; ppl: 339.95; xent: 5.83; lr: 0.00006; 48524/58874 tok/s; 364 sec
[2019-09-05 12:34:41,417 INFO] Step 500/50000; acc: 15.14; ppl: 235.72; xent: 5.46; lr: 0.00008; 46461/56365 tok/s; 439 sec
[2019-09-05 12:35:53,654 INFO] Step 600/50000; acc: 17.41; ppl: 176.58; xent: 5.17; lr: 0.00009; 48243/58731 tok/s; 511 sec
[2019-09-05 12:37:06,072 INFO] Step 700/50000; acc: 19.30; ppl: 139.99; xent: 4.94; lr: 0.00011; 48149/58553 tok/s; 584 sec

Commit 4a8c4ac
[2019-09-05 12:42:05,446 INFO] Step 100/50000; acc: 3.61; ppl: 6631.33; xent: 8.80; lr: 0.00002; 23708/28735 tok/s; 147 sec
[2019-09-05 12:43:16,428 INFO] Step 200/50000; acc: 7.97; ppl: 1824.98; xent: 7.51; lr: 0.00003; 48888/59853 tok/s; 218 sec
[2019-09-05 12:44:27,514 INFO] Step 300/50000; acc: 11.75; ppl: 524.17; xent: 6.26; lr: 0.00005; 48952/59220 tok/s; 289 sec
[2019-09-05 12:45:38,700 INFO] Step 400/50000; acc: 14.74; ppl: 278.04; xent: 5.63; lr: 0.00006; 49226/59553 tok/s; 360 sec
[2019-09-05 12:46:53,512 INFO] Step 500/50000; acc: 17.45; ppl: 181.13; xent: 5.20; lr: 0.00008; 46560/56793 tok/s; 435 sec
[2019-09-05 12:48:05,109 INFO] Step 600/50000; acc: 19.77; ppl: 132.15; xent: 4.88; lr: 0.00009; 48644/59191 tok/s; 507 sec
[2019-09-05 12:49:16,393 INFO] Step 700/50000; acc: 22.03; ppl: 101.56; xent: 4.62; lr: 0.00011; 48692/59330 tok/s; 578 sec

The accuracy / ppl seem much better on the old FusedAdam.

Any clue ?

@FDecaYed @mcarilli

Compare FusedAdam at commit 4a8c4ac with master after PR #424, using the reported OpenNMT-py configuration. Reproduce the CUDA out-of-memory and training-metric differences, then identify which changes in PR #424 account for the regression; done means explaining or correcting the memory and accuracy discrepancy.