Small language model to run on 1 TPU.

  Run this on 2x2 on languagemodel_lm1b32k_packed for 272000 steps (10 epochs)
  Results:
         params/10^9  log-ppl(per-token)
         0.14         3.202

  Returns:
    a hparams

def mtf_transformer_lm_baseline():
  """Small language model to run on 1 TPU.

  Run this on 2x2 on languagemodel_lm1b32k_packed for 272000 steps (10 epochs)
  Results:
         params/10^9  log-ppl(per-token)
         0.14         3.202

  Returns:
    a hparams
  """
  hparams = mtf_transformer_paper_lm(-1)
  hparams.batch_size = 128
  hparams.learning_rate_decay_steps = 27200  # one epoch on lm1b
  hparams.mesh_shape = "batch:8"
  return hparams