Introducing experimental gradient accumulation API

[rpsilva-aws]

In this PR, we introduce experimental.gradient_accumulation which leverages XLA's While op to accumulate gradients.

Training loop with traditional gradient accumulation
===> Preparing data..
Epoch 0 step 8 loss 1.1098170280456543
Epoch 0 step 16 loss 1.1719611883163452
Epoch 0 step 24 loss 3.453134536743164
Epoch 0 step 32 loss 2.518792152404785
Epoch 0 step 40 loss 6.67546272277832
Epoch 0 step 48 loss 4.609560012817383
Epoch 0 step 56 loss 5.953202247619629
Epoch 0 step 64 loss 1.325960636138916
Training loop with XLA's `While` gradient accumulation
===> Preparing data..
Epoch 0 step 8 loss 1.1098170280456543
Epoch 0 step 16 loss 1.1719611883163452
Epoch 0 step 24 loss 3.453134536743164
Epoch 0 step 32 loss 2.518792152404785
Epoch 0 step 40 loss 6.67546272277832
Epoch 0 step 48 loss 4.609560012817383
Epoch 0 step 56 loss 5.953202247619629
Epoch 0 step 64 loss 1.325960636138916

[rpsilva-aws]

@jeffhataws @tengyifei

[tengyifei]

@rpsilva-aws do you plan on merging this into r2.6?

[rpsilva-aws]

@tengyifei Ideally, yes. It's perfectly fine for the 3-layer MLP, but we're seeing a small difference for Llama runs (difference being, from a previous local patch set that was just before cleaning some of the code), so we're just quickly identifying what it is.

[tengyifei]

Okay, please aim to sort out all critical issues by Jan 21 if you're aiming for 2.6 so that we could review and cherrypick it by Jan 22. 2.6 release is quicking drawing in and I would like a few days to test all the builds.