Add llama 3 tokenizer

Author: sychen52

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-1B-v3-flash-single-host.txt

@@ -137,7 +137,7 @@ model.decoder.emb.token_emb.param_init.init_by_param_name['.*weight$'].scale: 1.
 model.decoder.emb.token_emb.param_init.klass: 'axlearn.common.param_init.DefaultInitializer'
 model.decoder.emb.token_emb.param_partition_spec[0]: None
 model.decoder.emb.token_emb.param_partition_spec[1]: 'model'
-model.decoder.eos_token_id: 1
+model.decoder.eos_token_id: 128001
 model.decoder.klass: 'axlearn.common.decoder.Decoder'
 model.decoder.logits_partition_spec[0][0]: 'data'
 model.decoder.logits_partition_spec[0][1]: 'expert'
@@ -148,7 +148,7 @@ model.decoder.output_dropout.klass: 'axlearn.common.layers.Dropout'
 model.decoder.output_norm.eps: 1e-05
 model.decoder.output_norm.forward_dtype: None
 model.decoder.output_norm.klass: 'axlearn.common.layers.RMSNorm'
-model.decoder.pad_token_id: 0
+model.decoder.pad_token_id: 128004
 model.decoder.transformer.klass: 'axlearn.common.attention.RepeatedTransformerLayer'
 model.decoder.transformer.layer.feed_forward.activation[0]: 'nn.silu'
 model.decoder.transformer.layer.feed_forward.activation[1]: 'linear'
@@ -267,7 +267,7 @@ model.decoder.transformer.num_layers: 16
 model.decoder.transformer.repeat.drop_output.fn: 'axlearn.common.repeat._drop_by_regex'
 model.decoder.transformer.repeat.drop_output.rules[0]: 'module_outputs.*'
 model.decoder.transformer.repeat.klass: 'axlearn.common.attention._TransformerRepeat'
-model.decoder.vocab_size: 128256
+model.decoder.vocab_size: 131072
 model.dtype: 'jax.numpy.float32'
 model.klass: 'axlearn.common.causal_lm.Model'
 model.param_init.init_by_param_name['.*weight$'].distribution: 'normal'

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-1B-v3-flash-single-host_init.txt

@@ -1,4 +1,4 @@
-decoder/emb/token_emb/weight: normal(0, 1.0 / fan_out), shape=[128256, 2048], axes=FanAxes(in_axis=-2, out_axis=-1, batch_axis=())
+decoder/emb/token_emb/weight: normal(0, 1.0 / fan_out), shape=[131072, 2048], axes=FanAxes(in_axis=-2, out_axis=-1, batch_axis=())
 decoder/transformer/repeat/layer/self_attention/norm/scale: constant(1.0)
 decoder/transformer/repeat/layer/self_attention/attention/i_proj/i_proj/qkv_proj/weight: normal(0, 1.0 / fan_in), shape=(2048, 48, 64), axes=FanAxes(in_axis=0, out_axis=(1, 2), batch_axis=())
 decoder/transformer/repeat/layer/self_attention/attention/o_proj/weight: normal(0, 1.0 / fan_in), shape=(2048, 32, 64), axes=FanAxes(in_axis=(1, 2), out_axis=0, batch_axis=())

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-1B-v3-flash-tiktoken-single-host.txt

@@ -0,0 +1,9 @@
+decoder/emb/token_emb/weight: normal(0, 1.0 / fan_out), shape=[128256, 2048], axes=FanAxes(in_axis=-2, out_axis=-1, batch_axis=())
+decoder/transformer/repeat/layer/self_attention/norm/scale: constant(1.0)
+decoder/transformer/repeat/layer/self_attention/attention/i_proj/i_proj/qkv_proj/weight: normal(0, 1.0 / fan_in), shape=(2048, 48, 64), axes=FanAxes(in_axis=0, out_axis=(1, 2), batch_axis=())
+decoder/transformer/repeat/layer/self_attention/attention/o_proj/weight: normal(0, 1.0 / fan_in), shape=(2048, 32, 64), axes=FanAxes(in_axis=(1, 2), out_axis=0, batch_axis=())
+decoder/transformer/repeat/layer/feed_forward/norm/scale: constant(1.0)
+decoder/transformer/repeat/layer/feed_forward/linear1_0/weight: normal(0, 1.0 / fan_in), shape=(2048, 8192), axes=FanAxes(in_axis=-2, out_axis=-1, batch_axis=())
+decoder/transformer/repeat/layer/feed_forward/linear1_1/weight: normal(0, 1.0 / fan_in), shape=(2048, 8192), axes=FanAxes(in_axis=-2, out_axis=-1, batch_axis=())
+decoder/transformer/repeat/layer/feed_forward/linear2/weight: normal(0, 1.0 / fan_in), shape=(8192, 2048), axes=FanAxes(in_axis=-2, out_axis=-1, batch_axis=())
+decoder/output_norm/scale: constant(1.0)

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-1B-v3-flash-tiktoken-single-host_init.txt

@@ -0,0 +1,10 @@
+====================weight_decay_scale root.optimizer====================
+decoder/emb/token_emb/weight: 1
+decoder/output_norm/scale: 1
+decoder/transformer/repeat/layer/feed_forward/linear1_0/weight: 1
+decoder/transformer/repeat/layer/feed_forward/linear1_1/weight: 1
+decoder/transformer/repeat/layer/feed_forward/linear2/weight: 1
+decoder/transformer/repeat/layer/feed_forward/norm/scale: 1
+decoder/transformer/repeat/layer/self_attention/attention/i_proj/i_proj/qkv_proj/weight: 1
+decoder/transformer/repeat/layer/self_attention/attention/o_proj/weight: 1
+decoder/transformer/repeat/layer/self_attention/norm/scale: 1

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-1B-v3-flash-tiktoken-single-host_regularizer.txt

@@ -0,0 +1,9 @@
+decoder/emb/token_emb/weight: normal(0, 1.0 / fan_out), shape=[128256, 2048], axes=FanAxes(in_axis=-2, out_axis=-1, batch_axis=())
+decoder/transformer/repeat/layer/self_attention/norm/scale: constant(1.0)
+decoder/transformer/repeat/layer/self_attention/attention/i_proj/i_proj/qkv_proj/weight: normal(0, 1.0 / fan_in), shape=(2048, 48, 64), axes=FanAxes(in_axis=0, out_axis=(1, 2), batch_axis=())
+decoder/transformer/repeat/layer/self_attention/attention/o_proj/weight: normal(0, 1.0 / fan_in), shape=(2048, 32, 64), axes=FanAxes(in_axis=(1, 2), out_axis=0, batch_axis=())
+decoder/transformer/repeat/layer/feed_forward/norm/scale: constant(1.0)
+decoder/transformer/repeat/layer/feed_forward/linear1_0/weight: normal(0, 1.0 / fan_in), shape=(2048, 8192), axes=FanAxes(in_axis=-2, out_axis=-1, batch_axis=())
+decoder/transformer/repeat/layer/feed_forward/linear1_1/weight: normal(0, 1.0 / fan_in), shape=(2048, 8192), axes=FanAxes(in_axis=-2, out_axis=-1, batch_axis=())
+decoder/transformer/repeat/layer/feed_forward/linear2/weight: normal(0, 1.0 / fan_in), shape=(8192, 2048), axes=FanAxes(in_axis=-2, out_axis=-1, batch_axis=())
+decoder/output_norm/scale: constant(1.0)

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-1B-v3-flash-tiktoken.txt

@@ -0,0 +1,10 @@
+====================weight_decay_scale root.optimizer====================
+decoder/emb/token_emb/weight: 1
+decoder/output_norm/scale: 1
+decoder/transformer/repeat/layer/feed_forward/linear1_0/weight: 1
+decoder/transformer/repeat/layer/feed_forward/linear1_1/weight: 1
+decoder/transformer/repeat/layer/feed_forward/linear2/weight: 1
+decoder/transformer/repeat/layer/feed_forward/norm/scale: 1
+decoder/transformer/repeat/layer/self_attention/attention/i_proj/i_proj/qkv_proj/weight: 1
+decoder/transformer/repeat/layer/self_attention/attention/o_proj/weight: 1
+decoder/transformer/repeat/layer/self_attention/norm/scale: 1

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-1B-v3-flash-tiktoken_init.txt

@@ -137,7 +137,7 @@ model.decoder.emb.token_emb.param_init.init_by_param_name['.*weight$'].scale: 1.
 model.decoder.emb.token_emb.param_init.klass: 'axlearn.common.param_init.DefaultInitializer'
 model.decoder.emb.token_emb.param_partition_spec[0]: None
 model.decoder.emb.token_emb.param_partition_spec[1]: 'model'
-model.decoder.eos_token_id: 1
+model.decoder.eos_token_id: 128001
 model.decoder.klass: 'axlearn.common.decoder.Decoder'
 model.decoder.logits_partition_spec[0][0]: 'data'
 model.decoder.logits_partition_spec[0][1]: 'expert'
@@ -148,7 +148,7 @@ model.decoder.output_dropout.klass: 'axlearn.common.layers.Dropout'
 model.decoder.output_norm.eps: 1e-05
 model.decoder.output_norm.forward_dtype: None
 model.decoder.output_norm.klass: 'axlearn.common.layers.RMSNorm'
-model.decoder.pad_token_id: 0
+model.decoder.pad_token_id: 128004
 model.decoder.transformer.klass: 'axlearn.common.attention.RepeatedTransformerLayer'
 model.decoder.transformer.layer.feed_forward.activation[0]: 'nn.silu'
 model.decoder.transformer.layer.feed_forward.activation[1]: 'linear'
@@ -267,7 +267,7 @@ model.decoder.transformer.num_layers: 16
 model.decoder.transformer.repeat.drop_output.fn: 'axlearn.common.repeat._drop_by_regex'
 model.decoder.transformer.repeat.drop_output.rules[0]: 'module_outputs.*'
 model.decoder.transformer.repeat.klass: 'axlearn.common.attention._TransformerRepeat'
-model.decoder.vocab_size: 128256
+model.decoder.vocab_size: 131072
 model.dtype: 'jax.numpy.float32'
 model.klass: 'axlearn.common.causal_lm.Model'
 model.param_init.init_by_param_name['.*weight$'].distribution: 'normal'

axlearn/experiments/testdata/axlearn.experiments.text.gpt.c4_trainer/fuji-1B-v3-flash-tiktoken_regularizer.txt

@@ -1,4 +1,4 @@
-decoder/emb/token_emb/weight: normal(0, 1.0 / fan_out), shape=[128256, 2048], axes=FanAxes(in_axis=-2, out_axis=-1, batch_axis=())
+decoder/emb/token_emb/weight: normal(0, 1.0 / fan_out), shape=[131072, 2048], axes=FanAxes(in_axis=-2, out_axis=-1, batch_axis=())
 decoder/transformer/repeat/layer/self_attention/norm/scale: constant(1.0)
 decoder/transformer/repeat/layer/self_attention/attention/i_proj/i_proj/qkv_proj/weight: normal(0, 1.0 / fan_in), shape=(2048, 48, 64), axes=FanAxes(in_axis=0, out_axis=(1, 2), batch_axis=())
 decoder/transformer/repeat/layer/self_attention/attention/o_proj/weight: normal(0, 1.0 / fan_in), shape=(2048, 32, 64), axes=FanAxes(in_axis=(1, 2), out_axis=0, batch_axis=())