code: def xmoe2_dense(sz):
    
    hparams = mtf_transformer.mtf_transformer_paper_lm(sz)
    hparams.attention_dropout = 0.0
    hparams.relu_dropout = 0.0
    hparams.layer_prepostprocess_dropout = 0.0
    hparams.max_length = 1024
    hparams.batch_size = 128
    hparams.learning_rate_schedule = "rsqrt_decay*linear_decay"
    hparams.learning_rate_decay_steps = 65536
    hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model"
    hparams.mesh_shape = "batch:32"
    return hparams
docs: Series of architectural experiments on language modeling.
  
    Larger models than the ones above.
  
    All models are trained on sequences of 1024 tokens.
  
    We assume infinite training data, so no dropout necessary.
    We process 2^36 tokens in training = 524288 steps at batch size 128
  
    TODO(noam): find a large enough dataset for these experiments.
  
    You can use languagemodel_wiki_noref_v32k_l1k, but this is too small,
    (1 epoch = ~46000 steps) so training will cover about 11 epochs.
  
    Note: configurations and code are likely to change without notice.
  
    Run on TPU 4x4 for 524288 steps unless otherwise indicated.
  
    Args:
      sz: an integer
  
    Returns:
      a hparams