Series of architectures for language modeling.

  We assume infinite training data, so no dropout necessary.

  You can use languagemodel_wiki_noref_v32k_l1k.
  (1 epoch = ~46000 steps).
  TODO(noam): find a large enough dataset for these experiments.

  Args:
    sz: an integer

  Returns:
    a hparams

def mtr_lm_dense(sz):
  """Series of architectures for language modeling.

  We assume infinite training data, so no dropout necessary.

  You can use languagemodel_wiki_noref_v32k_l1k.
  (1 epoch = ~46000 steps).
  TODO(noam): find a large enough dataset for these experiments.

  Args:
    sz: an integer

  Returns:
    a hparams
  """
  n = 2 ** sz
  hparams = mtf_unitransformer_base()
  hparams.d_model = 1024
  hparams.max_length = 1024
  hparams.batch_size = 128
  # Parameters for my_layer_stack()
  hparams.num_hidden_layers = 6
  hparams.d_ff = 8192 * n
  hparams.d_kv = 256
  hparams.num_heads = 8 * n
  hparams.learning_rate_decay_steps = 65536
  hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model"
  hparams.mesh_shape = "batch:32"
  return hparams