Add llama 3 tokenizer #850
Conversation
| @@ -67,7 +67,7 @@ class Version(enum.Enum): | |||
| VOCAB_SIZE = { | |||
| Version.V1: 32 * 1024, | |||
| Version.V2: 32 * 1024, | |||
| Version.V3: 128256, | |||
| Version.V3: 128 * 1024, | |||
There was a problem hiding this comment.
llama 3 is supposed to be 128256. Are there any plans to stick to standard llama 3?
There was a problem hiding this comment.
I add -tiktoken configs and they will use 128256.
the configs without -tiktoken will stay unchanged. vocab_size is still 128*1024 and using bpe tokenizer.
There was a problem hiding this comment.
ah nice! I didn't catch that part.
There was a problem hiding this comment.
Thanks! Could you confirm with @kelvin-zou that the new configs still work on TPUs and GPUs?
| @@ -301,6 +305,10 @@ def model_config( | |||
| lm_head=lm_head_cfg, | |||
| dropout_rate=dropout_rate, | |||
| ) | |||
| if pad_token_id: | |||
There was a problem hiding this comment.
Since pad_token_id=0 is a valid token:
| if pad_token_id: | |
| if pad_token_id is not None: |
| @@ -301,6 +305,10 @@ def model_config( | |||
| lm_head=lm_head_cfg, | |||
| dropout_rate=dropout_rate, | |||
| ) | |||
| if pad_token_id: | |||
| decoder_cfg.set(pad_token_id=pad_token_id) | |||
| if eos_token_id: | |||
There was a problem hiding this comment.
| if eos_token_id: | |
| if eos_token_id is not None: |
| 128256: "bpe_128k_c4.model", | ||
| } | ||
|
|
||
| def _vocab_cfg(size: int): |
There was a problem hiding this comment.
nit -
| def _vocab_cfg(size: int): | |
| def _vocab_cfg(vocab_size: int): |
| if size == 128256: | ||
| # TikToken. | ||
| return config_for_class(FujiV3Vocabulary).set(filename="Llama-3-tokenizer.json") | ||
| raise ValueError(f"size {size} tokenizer does not exist.") |
There was a problem hiding this comment.
| raise ValueError(f"size {size} tokenizer does not exist.") | |
| raise ValueError(f"Tokenizer with vocab size {size} does not exist.") |
| pad_token_id: Optional[int] = None, | ||
| eos_token_id: Optional[int] = None, |
There was a problem hiding this comment.
Please update docstring accordingly.
axlearn/experiments/text/gpt/fuji.py
Outdated
| @@ -423,6 +421,9 @@ def get_trainer_kwargs( | |||
| raise NotImplementedError(f"Unknown model size {model_size}.") | |||
| model_kwargs = trainer_kwargs.pop("model_kwargs") | |||
| model_kwargs.setdefault("vocab_size", vocab_size) | |||
| if version == Version.V3 and vocab_size == 128256: # tiktoken tokenizer | |||
There was a problem hiding this comment.
Can we use a private const for the tiktoken size?
_TIKTOKEN_VOCAB_SIZE = 128256
There was a problem hiding this comment.
I will add a new version called V3_tiktoken. I think it will make things cleaner.
axlearn/experiments/text/gpt/fuji.py
Outdated
| vocab_size = VOCAB_SIZE[version] | ||
| if tiktoken: | ||
| suffix += "-tiktoken" | ||
| vocab_size = 128256 |
| @@ -27,6 +27,65 @@ | |||
| # Use cpu for the test. | |||
| jax.config.update("jax_platform_name", "cpu") | |||
|
|
|||
| config_dict_1b = { | |||
There was a problem hiding this comment.
Add a link to the json config as a reference?
| filename = os.path.join(data_dir, "tokenizers", "hf", filename) | ||
| if filename.startswith("gs:") or filename.startswith("s3:"): | ||
| # Create a different file for each usage. | ||
| tmp = tempfile.mkdtemp() |
There was a problem hiding this comment.
Should we cleanup tmp after tokenizer has been initialized?
| 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 |
There was a problem hiding this comment.
@samos123 for example here it is using tiktoken and vocab_size is 128256.
| @@ -301,6 +305,10 @@ def model_config( | |||
| lm_head=lm_head_cfg, | |||
| dropout_rate=dropout_rate, | |||
| ) | |||
| if pad_token_id: | |||
| @@ -301,6 +305,10 @@ def model_config( | |||
| lm_head=lm_head_cfg, | |||
| dropout_rate=dropout_rate, | |||
| ) | |||
| if pad_token_id: | |||
| decoder_cfg.set(pad_token_id=pad_token_id) | |||
| if eos_token_id: | |||
| 128256: "bpe_128k_c4.model", | ||
| } | ||
|
|
||
| def _vocab_cfg(size: int): |
| if size == 128256: | ||
| # TikToken. | ||
| return config_for_class(FujiV3Vocabulary).set(filename="Llama-3-tokenizer.json") | ||
| raise ValueError(f"size {size} tokenizer does not exist.") |
| pad_token_id: Optional[int] = None, | ||
| eos_token_id: Optional[int] = None, |
axlearn/experiments/text/gpt/fuji.py
Outdated
| @@ -423,6 +421,9 @@ def get_trainer_kwargs( | |||
| raise NotImplementedError(f"Unknown model size {model_size}.") | |||
| model_kwargs = trainer_kwargs.pop("model_kwargs") | |||
| model_kwargs.setdefault("vocab_size", vocab_size) | |||
| if version == Version.V3 and vocab_size == 128256: # tiktoken tokenizer | |||
There was a problem hiding this comment.
I will add a new version called V3_tiktoken. I think it will make things cleaner.
axlearn/experiments/text/gpt/fuji.py
Outdated
| vocab_size = VOCAB_SIZE[version] | ||
| if tiktoken: | ||
| suffix += "-tiktoken" | ||
| vocab_size = 128256 |
| @@ -27,6 +27,65 @@ | |||
| # Use cpu for the test. | |||
| jax.config.update("jax_platform_name", "cpu") | |||
|
|
|||
| config_dict_1b = { | |||
| filename = os.path.join(data_dir, "tokenizers", "hf", filename) | ||
| if filename.startswith("gs:") or filename.startswith("s3:"): | ||
| # Create a different file for each usage. | ||
| tmp = tempfile.mkdtemp() |
sychen52
force-pushed
the
llama
branch
2 times, most recently
from
3e8d4f9 to
81f3e7b
Compare
add a new version called V3_TIKTOKEN. other edits based on suggestions.
No description provided.