Set of hyperparameters.

def mtf_transformer_base():
  """Set of hyperparameters."""
  hparams = common_hparams.basic_params1()
  hparams.no_data_parallelism = True
  hparams.use_fixed_batch_size = True
  hparams.add_hparam("mtf_mode", True)
  hparams.batch_size = 64
  hparams.max_length = 256
  hparams.add_hparam("d_model", 512)
  hparams.add_hparam("d_kv", 128)
  hparams.add_hparam("local_attention_window_size", 128)
  hparams.label_smoothing = 0.1
  # 8-way model-parallelism
  hparams.add_hparam("mesh_shape", "model:8")
  hparams.add_hparam("layout", "batch:batch;vocab:model;d_ff:model;heads:model")
  hparams.add_hparam("num_heads", 8)
  hparams.add_hparam("d_ff", 2048)
  hparams.add_hparam("encoder_replicate_factor", 1)
  hparams.add_hparam("decoder_replicate_factor", 1)
  hparams.add_hparam("encoder_layers", ["att", "drd"] * 6)
  hparams.add_hparam("decoder_layers", ["att", "enc_att", "drd"] * 6)
  hparams.add_hparam("attention_dropout", 0.1)
  hparams.add_hparam("relu_dropout", 0.1)
  hparams.layer_prepostprocess_dropout = 0.1

  # Describes what model architecture:
  #   "encdec": encoder + autoregressive decoder
  #   "decoder": single-stack autoregressive sequence model.
  #   "encoder": single-stack non-autoregressive model
  #      with equal-length inputs and outputs.
  hparams.add_hparam("transformer_type", "encdec")

  # What does the decoder do:
  #   "autoregressive": Decoder left to right
  #   "denoising": Fills in masked-out values simultaneously
  hparams.add_hparam("decoder_type", "autoregressive")

  # Parameters describing the noising algorithm for denoising decoders
  hparams.add_hparam("noising_spec_train", {"type": "mask", "prob": 0.15})
  hparams.add_hparam("noising_spec_eval", {"type": "mask", "prob": 0.15})
  # during training, we use the eval noiser with this probability
  hparams.add_hparam("noising_use_eval_during_train", 0.1)

  # round up vocab sizes to be a multiple of this value
  hparams.vocab_divisor = 128

  # options are dense_relu_dense, moe, hmoe
  hparams.add_hparam("feedforward_layer", "drd")

  # If True, then reuse targets_embedding_var * rsqrt(d_model) as softmax_var
  # If hparams.transformer_type == "encoder", then there is no targets embedding
  # so we reuse the inputs embedding instead.
  hparams.shared_embedding_and_softmax_weights = True
  # Reuse targets_embedding_var as inputs_embedding_var
  # relevant only if hparams.transformer_type == "encdec"
  hparams.shared_embedding = True
  hparams.optimizer = "Adafactor"
  hparams.learning_rate_schedule = "linear_warmup*rsqrt_decay*linear_decay"
  hparams.learning_rate_warmup_steps = 10000
  hparams.add_hparam("master_dtype", "bfloat16")
  hparams.add_hparam("slice_dtype", "float32")
  hparams.activation_dtype = "bfloat16"

  # These parameters make Transformer model compatible with MtfTransformer
  # Do not override these, as mtf_transformer does not support other options.
  hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
  hparams.bottom = {
      "inputs": modalities.identity_bottom,
      "targets": modalities.identity_bottom,
  }
  hparams.top = {
      "targets": modalities.identity_top,
  }

  # Parameters for computing the maximum decode length in beam search.
  # Maximum decode length is:
  #    min(max_length,
  #        decode_length_multiplier * input_length + decode_length_constant)
  hparams.add_hparam("decode_length_multiplier", 1.5)
  hparams.add_hparam("decode_length_constant", 10.0)

  # If nonzero, we split the batch across two tensor-dimensions named
  # "outer_batch" and "inner_batch", allowing for splitting across two mesh
  # dimensions.  This is necessary for hierarchical mixture of experts.
  # The two tensor dimensions have sizes hparams.outer_batch_size and
  # hparams.batch_size // hparams.outer_batch_size.
  hparams.add_hparam("outer_batch_size", 0)

  # TODO(noam): file a bug
  hparams.add_hparam("reshape_logits_hack", False)
  hparams.add_hparam("compression_factor", 4)

  return hparams