Model incorporating mixture-of-experts, local and global attention.

  ~6B parameters

  32 experts in 3 hierarchichal moe layers.

  Returns:
    a hparams

def mtr_lm_v1():
  """Model incorporating mixture-of-experts, local and global attention.

  ~6B parameters

  32 experts in 3 hierarchichal moe layers.

  Returns:
    a hparams
  """
  hparams = mtr_lm_dense(0)
  hparams.layers = (["local_self_att", "local_self_att", "drd",
                     "self_att", "drd", "local_self_att",
                     "local_self_att", "moe_2d"] * 4)[:-1]
  hparams.d_kv = 128
  hparams.moe_expert_x = 8
  hparams.moe_expert_y = 4
  hparams.moe_hidden_size = 32768
  hparams.d_ff = 2048
  hparams.num_memory_heads = 0
  hparams.mesh_shape = "b0:4;b1:8"
  hparams.layout = "outer_batch:b0;inner_batch:b1,expert_x:b1,expert_y:b0"
  hparams.outer_batch_size = 4
  return hparams