code: def wiki_2x2_base():
    
    hparams = mtf_transformer.mtf_transformer_base_lm()
    hparams.shared_embedding_and_softmax_weights = False
    # no dropout - dataset is big enough to avoid overfitting.
    hparams.attention_dropout = 0.0
    hparams.relu_dropout = 0.0
    hparams.layer_prepostprocess_dropout = 0.0
    hparams.max_length = 1024
    # 4 sequences per core
    hparams.batch_size = 32
    # We don't use linear decay in these experiments, since we don't want
    # a sharp jump in quality at the end of the training schedule.
    # You can insert this once you find the right architecture.
    hparams.learning_rate_schedule = "rsqrt_decay"
    hparams.mesh_shape = "all:8"
    hparams.layout = "batch:all;experts:all"
  
    # parameters for mixture-of-experts
    moe.set_default_moe_hparams(hparams)
    hparams.moe_num_experts = 16
    hparams.moe_hidden_size = 8192
  
    hparams.decoder_layers = ["att", "drd"] * 6
    hparams.d_model = 1024
    hparams.d_ff = 2048
    hparams.d_kv = 128
    hparams.num_heads = 4
  
    return hparams
docs: Set of architectural experiments - language model on wikipedia on a 2x2.
  
    1 epoch = ~180k steps at batch size 32 - we may never finish an epoch!
  
    Returns:
      a hparams