def xmoe_wiki_base(): """Series of architectural experiments on wikipedia text. For all of these architectures, we run on languagemodel_wiki_noref_v8k_l1k for 3 epochs. (training set has ~7390100 sequences each of length 1024) 1 epoch = 115000 steps at batch_size=64 Results: model params(M) einsum alltoall mxu-util log-ppl(1ep) (3ep) Note: configurations and code are likely to change without notice. Returns: a hparams """ hparams = mtf_transformer.mtf_transformer_base() # The following hparams are constant across all these experiments. hparams.label_smoothing = 0.0 hparams.max_length = 1024 hparams.batch_size = 64 hparams.d_model = 1024 hparams.d_kv = 128 hparams.num_heads = 8 hparams.shared_embedding_and_softmax_weights = False hparams.learning_rate_decay_steps = 115000 # We will vary the following parameters related to the ffn/moe layers. hparams.feedforward_layer = "dense_relu_dense" hparams.d_ff = 8192 hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model" hparams.mesh_shape = "batch:32" return hparams
def xmoe_dense_4k(): """Series of architectural experiments on cheap language models. For all of these architectures, we run on languagemodel_lm1b8k_packed for 32k-96 steps (1-3 epochs) on one TPU (8 cores). All log-perplexities are per-token - multiply by 1.298 for per-word Results: model params(M) einsum alltoall mxu-util log-ppl(1ep) (3ep) xmoe_dense_4k 30 3.0e12 0 45% 3.31 xmoe_dense_8k 46 4.7e12 0 49% 3.24 xmoe_dense_64k 282 2.8e13 0 3.06 xmoe_top_2 282 4.0e12 3.4e8 36% 3.07 xmoe_top_2_c15 282 4.5e12 4.0e8 38% 3.07 xmoe_2d 282 5.3e12 7.6e8 34% 3.06 Trained at 4x the batch size: xmoe_2d_88 1090 2.1e13 3.0e9 24% Note: configurations and code are likely to change without notice. Returns: a hparams """ hparams = mtf_transformer.mtf_transformer_base() # The following hparams are constant across all these experiments. hparams.label_smoothing = 0.0 hparams.batch_size = 128 hparams.d_model = 512 hparams.d_kv = 128 hparams.num_heads = 4 hparams.num_decoder_layers = 4 hparams.shared_embedding_and_softmax_weights = False hparams.learning_rate_schedule = "rsqrt_decay" # We will vary the following parameters related to the ffn/moe layers. hparams.feedforward_layer = "dense_relu_dense" hparams.d_ff = 4096 hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model" hparams.mesh_shape = "batch:8" return hparams
def xmoe_dense_4k(): """Small transformer language model.""" hparams = mtf_transformer.mtf_transformer_base() # The following hparams are constant across all these experiments. hparams.label_smoothing = 0.0 hparams.batch_size = 128 hparams.d_model = 512 hparams.d_kv = 128 hparams.num_heads = 4 hparams.num_decoder_layers = 4 hparams.shared_embedding_and_softmax_weights = False hparams.learning_rate_schedule = "rsqrt_decay" # We will vary the following parameters related to the ffn/moe layers. hparams.feedforward_layer = "dense_relu_dense" hparams.d_ff = 4096 hparams.moe_num_experts = 16 hparams.moe_overhead_train = 1.0 hparams.moe_overhead_eval = 2.0 hparams.moe_loss_coef = 1e-3 hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model" hparams.mesh_shape = "batch:8" return hparams