GSPMD "2D Partitioning" Training Speed

[agemagician]

Hello,

I am testing the training speed using 2D Partitioning, and it gives me approximately 9.5 seconds per step, while on the paper, it should be approximately 6.5 giving a similar model size.

I am testing it with "scalable_t5" while using the following related parameters:

network.Transformer.config = @network.T5Config()
network.T5Config:
  scan_layers = True
  remat_policy = 'minimal'

partitioning.PjitPartitioner:
  num_partitions = 4

partitioning.standard_logical_axis_rules:
  activation_partitioning_dims = 2
  parameter_partitioning_dims = 2

trainer.Trainer:
  num_microbatches = 4 

BATCH_SIZE = 1024

Paper "Table 2":
https://arxiv.org/pdf/2105.04663.pdf

@adarob Is it possible to share the configuration that achieved such a speed on TPU V3 ?