How to set batch size of training?

[desperadoola]

When I try to change the batch size using --gin_param="sequences_per_batch=128" or --gin_param="tokens_per_batch=65536", the batch size seems always to be 32?

INFO:tensorflow:serialize_num_microbatches: tokens_per_microbatch_per_replica=2048 batch_dim=Dimension(name='batch', size=32) sequence_length={'inputs': 512, 'targets': 114} batch_per_replica=4 num_microbatches=1 I1208 11:05:22.407459 140391696871040 utils.py:1440] serialize_num_microbatches: tokens_per_microbatch_per_replica=2048 batch_dim=Dimension(name='batch', size=32) sequence_length={'inputs': 512, 'targets': 114} batch_per_replica=4 num_microbatches=1

[desperadoola]

I successfully set it using --gin_param="utils.run.batch_size=('tokens_per_batch', 65536)"

[craffel]

Good work!

[desperadoola]

Is there any instruction on how to set tokens_per_microbatch_per_replica ?

[craffel]

You should only need to set tokens_per_microbatch_per_replica to something other than None if you want to use a batch size which is too large to fit in memory. Our training code will automatically split up too-large batches into microbatches and accumulate gradients so that the full batch size is computed.

@nshazeer FYI

[desperadoola]

Thanks, but still, I don't figure out serialize_num_microbatches and mtf.tensor_dim_to_size_per_split in mesh_tensorflow transformer.

Is the default value tokens_per_microbatch_per_replica=2048 ok for different settings, for example, when we change the model size and use different TPU?