Using Gradient Accumulation

Hey guys!

I am about to pretrain a monolingual model using T5X (thank you for this!).

The routine I'll be following is based on ByT5 paper. However, I currently have access to a smaller TPU (v3-8 core), so 2²⁰ would not fit into its memory. To accomplish similar results, I am thinking of using gradient accumulation, so I can emulate the same batch size used for pertaining the original ByT5.

I couldn't find any documentation about this, but looking through the code, I guess I would have to:

1. Specify BATCH_SIZE = 1024 (2²⁰/1024)
2. Specify trainer.num_microbatches = 16

With this, I hope to fit 64 examples in a step (does T5X auto distribute 8 samples per core here?), but update gradients every 16 steps, emulating a 1024 batch size.

The resulting gin would be something like:

```gin
include 't5x/examples/t5/byt5/small.gin'
include 't5x/configs/runs/pretrain.gin'

TASK_FEATURE_LENGTHS = {"inputs": 1024, "targets": 189}
TRAIN_STEPS = 1_000_000
DROPOUT_RATE = 0.0
BATCH_SIZE = 1024

trainer.Trainer:
num_microbatches = 16
```

Is that correct?

Start with trainer.Trainer and the pretrain.gin and ByT5 small.gin configuration snippets named in the issue. Verify how num_microbatches and BATCH_SIZE interact on the described TPU setup; done means a clear documented answer that confirms or corrects the proposed configuration.

Written by the indexing model from the issue text.