def mtf_transformer2_base():
  """Hyperparameters common to both unitransformer and bitransformer."""
  hparams = common_hparams.basic_params1()

  hparams.add_hparam("d_model", 1024)
  hparams.batch_size = 4
  hparams.max_length = 1024
  hparams.label_smoothing = 0.0
  # a small positive value - this seems important for stability when training
  # with bfloat16 activations.
  hparams.add_hparam("z_loss", 1e-4)

  # hparams applying to both encoder and decoder layer stacks.
  hparams.add_hparam("d_ff", 2048)
  hparams.add_hparam("d_kv", 128)
  hparams.add_hparam("attention_dropout", 0.0)
  hparams.add_hparam("relu_dropout", 0.0)
  hparams.del_hparam("num_heads")
  hparams.del_hparam("num_hidden_layers")
  hparams.layer_prepostprocess_dropout = 0.0
  hparams.add_hparam("extra_logit", False)
  # number of experts for moe_1d
  hparams.moe_num_experts = 32
  # number of experts for moe_2d = moe_expert_x * moe_expert_y
  hparams.add_hparam("moe_expert_x", 8)
  hparams.add_hparam("moe_expert_y", 4)
  hparams.add_hparam("moe_hidden_size", 32768)

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

  hparams.optimizer = "Adafactor"
  hparams.learning_rate_schedule = "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"

  # 8-way model-parallelism
  hparams.add_hparam("mesh_shape", "model:8")
  hparams.add_hparam("layout", "batch:batch;vocab:model;d_ff:model;heads:model")

  # 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)

  hparams.shared_embedding_and_softmax_weights = False
  # length for training or decoding - defaults to max_length
  hparams.add_hparam("length", 0)

  # These parameters make Transformer model compatible with mtf
  # Do not override these.
  hparams.no_data_parallelism = True
  hparams.use_fixed_batch_size = True
  hparams.add_hparam("mtf_mode", True)
  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,
  }
  hparams.add_hparam("beam_size", 1)

  # If this is True, then in a packed dataset (where exaples are concatenated
  # to form longer examples) we use the global position (within the concatenated
  # sequence) to compute the positional embedding, instead of the position
  # within the individual sequence.  This is counterintuitive, but for some
  # reason, it keeps the model from diverging.
  hparams.add_hparam("use_global_position_in_packed_sequence", True)

  return hparams
Hyperparameters common to both unitransformer and bitransformer.
tensor2tensor/models/mtf_transformer2.py