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
def mtf_transformer_base ( ) : 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 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 hparams . add_hparam ( "transformer_type" , "encdec" ) hparams . add_hparam ( "decoder_type" , "autoregressive" ) hparams . add_hparam ( "noising_spec_train" , { "type" : "mask" , "prob" : 0.15 } ) hparams . add_hparam ( "noising_spec_eval" , { "type" : "mask" , "prob" : 0.15 } ) hparams . add_hparam ( "noising_use_eval_during_train" , 0.1 ) hparams . vocab_divisor = 128 hparams . add_hparam ( "feedforward_layer" , "drd" ) hparams . shared_embedding_and_softmax_weights = True 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" hparams . clip_grad_norm = 0. hparams . bottom = { "inputs" : modalities . identity_bottom , "targets" : modalities . identity_bottom , } hparams . top = { "targets" : modalities . identity_top , } hparams . add_hparam ( "decode_length_multiplier" , 1.5 ) hparams . add_hparam ( "decode_length_constant" , 10.0 ) hparams . add_hparam ( "outer_batch_size" , 0 ) hparams . add_hparam ( "reshape_logits_hack" , False ) hparams . add_hparam ( "compression_factor" , 4 ) return hparams
Set of hyperparameters.