Add llama 3 tokenizer (apple#850)

* Add llama 3 tokenizer

add a new version called V3_TIKTOKEN.

other edits based on suggestions.

* Handle special tokens like other vocabularies.

* use encode instead of encode_batch

Author: sychen52

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

@@ -264,7 +264,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.txt

@@ -264,7 +264,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_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-single-host.txt

@@ -231,7 +231,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-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-tiktoken-flash-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-tiktoken-flash-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