| |
| |
| |
| |
| @@ -17,10 +17,10 @@ import json |
| import os |
| import shutil |
| import warnings |
| - |
| +from typing import List |
| import torch |
| |
| -from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer, PreTrainedTokenizerFast |
| +from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer, PreTrainedTokenizerFast, GenerationConfig |
| from transformers.convert_slow_tokenizer import TikTokenConverter |
| |
| |
| @@ -85,8 +85,12 @@ NUM_SHARDS = { |
| "65B": 8, |
| "70B": 8, |
| "70Bf": 8, |
| + "405B": 8, |
| + "405B-MP16": 16, |
| } |
| |
| +CONTEXT_LENGTH_FOR_VERSION = {"3.1": 131072, "3": 8192, "2": 4096, "1": 2048} |
| + |
| |
| def compute_intermediate_size(n, ffn_dim_multiplier=1, multiple_of=256): |
| return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3)) + multiple_of - 1) // multiple_of) |
| @@ -107,9 +111,10 @@ def write_model( |
| input_base_path, |
| model_size=None, |
| safe_serialization=True, |
| - llama_version=1, |
| + llama_version="1", |
| vocab_size=None, |
| num_shards=None, |
| + instruct=False, |
| ): |
| os.makedirs(model_path, exist_ok=True) |
| tmp_model_path = os.path.join(model_path, "tmp") |
| @@ -125,18 +130,11 @@ def write_model( |
| dims_per_head = dim // n_heads |
| base = params.get("rope_theta", 10000.0) |
| inv_freq = 1.0 / (base ** (torch.arange(0, dims_per_head, 2).float() / dims_per_head)) |
| - if base > 10000.0 and llama_version != 3: |
| + if base > 10000.0 and float(llama_version) < 3: |
| max_position_embeddings = 16384 |
| else: |
| - # Depending on the Llama version, the default max_position_embeddings has different values. |
| - if llama_version == 1: |
| - max_position_embeddings = 2048 |
| - elif llama_version == 2: |
| - max_position_embeddings = 4096 |
| - elif llama_version == 3: |
| - max_position_embeddings = 8192 |
| - |
| - vocab_size = vocab_size if vocab_size is not None else 32000 |
| + max_position_embeddings = CONTEXT_LENGTH_FOR_VERSION[llama_version] |
| + |
| if params.get("n_kv_heads", None) is not None: |
| num_key_value_heads = params["n_kv_heads"] # for GQA / MQA |
| num_key_value_heads_per_shard = num_key_value_heads // num_shards |
| @@ -144,8 +142,7 @@ def write_model( |
| else: # compatibility with other checkpoints |
| num_key_value_heads = n_heads |
| num_key_value_heads_per_shard = n_heads_per_shard |
| - key_value_dim = dims_per_head * num_key_value_heads |
| - print(num_shards, num_key_value_heads, num_key_value_heads_per_shard, key_value_dim) |
| + key_value_dim = dim |
| |
| # permute for sliced rotary |
| def permute(w, n_heads, dim1=dim, dim2=dim): |
| @@ -159,11 +156,9 @@ def write_model( |
| loaded = torch.load(os.path.join(input_base_path, "consolidated.00.pth"), map_location="cpu") |
| else: |
| # Sharded |
| - loaded = [ |
| - torch.load(os.path.join(input_base_path, file), map_location="cpu") |
| - for file in os.listdir(input_base_path) |
| - if file.endswith(".pth") |
| - ] |
| + checkpoint_list = sorted([file for file in os.listdir(input_base_path) if file.endswith(".pth")]) |
| + print("Loading in order:", checkpoint_list) |
| + loaded = [torch.load(os.path.join(input_base_path, file), map_location="cpu") for file in checkpoint_list] |
| param_count = 0 |
| index_dict = {"weight_map": {}} |
| for layer_i in range(n_layers): |
| @@ -263,7 +258,7 @@ def write_model( |
| "lm_head.weight": loaded["output.weight"], |
| } |
| else: |
| - concat_dim = 0 if llama_version == 3 else 1 |
| + concat_dim = 0 if llama_version in ['3', '3.1'] else 1 |
| state_dict = { |
| "model.norm.weight": loaded[0]["norm.weight"], |
| "model.embed_tokens.weight": torch.cat( |
| @@ -282,6 +277,18 @@ def write_model( |
| write_json(index_dict, os.path.join(tmp_model_path, "pytorch_model.bin.index.json")) |
| ffn_dim_multiplier = params["ffn_dim_multiplier"] if "ffn_dim_multiplier" in params else 1 |
| multiple_of = params["multiple_of"] if "multiple_of" in params else 256 |
| + |
| + if llama_version in ['3', '3.1']: |
| + bos_token_id = 128000 |
| + |
| + if instruct: |
| + eos_token_id = [128001, 128008, 128009] |
| + else: |
| + eos_token_id = 128001 |
| + else: |
| + bos_token_id = 1 |
| + eos_token_id = 2 |
| + |
| config = LlamaConfig( |
| hidden_size=dim, |
| intermediate_size=compute_intermediate_size(dim, ffn_dim_multiplier, multiple_of), |
| @@ -292,11 +299,21 @@ def write_model( |
| vocab_size=vocab_size, |
| rope_theta=base, |
| max_position_embeddings=max_position_embeddings, |
| - bos_token_id=128000 if llama_version == 3 else 1, |
| - eos_token_id=128001 if llama_version == 3 else 2, |
| + bos_token_id=bos_token_id, |
| + eos_token_id=eos_token_id, |
| ) |
| config.save_pretrained(tmp_model_path) |
| |
| + if instruct: |
| + generation_config = GenerationConfig( |
| + do_sample=True, |
| + temperature=0.6, |
| + top_p=0.9, |
| + bos_token_id=bos_token_id, |
| + eos_token_id=eos_token_id, |
| + ) |
| + generation_config.save_pretrained(tmp_model_path) |
| + |
| # Make space so we can load the model properly now. |
| del state_dict |
| del loaded |
| @@ -313,7 +330,7 @@ def write_model( |
| |
| |
| class Llama3Converter(TikTokenConverter): |
| - def __init__(self, vocab_file, num_reserved_special_tokens=256, **kwargs): |
| + def __init__(self, vocab_file, special_tokens=None, instruct=False, model_max_length=None, **kwargs): |
| super().__init__(vocab_file, **kwargs) |
| tokenizer = self.converted() |
| chat_template = ( |
| @@ -327,34 +344,27 @@ class Llama3Converter(TikTokenConverter): |
| "{% endfor %}" |
| "{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}" |
| ) |
| - num_reserved_special_tokens = 256 |
| - special_tokens = [ |
| - "<|begin_of_text|>", |
| - "<|end_of_text|>", |
| - "<|reserved_special_token_0|>", |
| - "<|reserved_special_token_1|>", |
| - "<|reserved_special_token_2|>", |
| - "<|reserved_special_token_3|>", |
| - "<|start_header_id|>", |
| - "<|end_header_id|>", |
| - "<|reserved_special_token_4|>", |
| - "<|eot_id|>", # end of turn |
| - ] + [f"<|reserved_special_token_{i}|>" for i in range(5, num_reserved_special_tokens - 5)] |
| tokenizer.add_special_tokens(special_tokens) |
| |
| self.tokenizer = PreTrainedTokenizerFast( |
| tokenizer_object=tokenizer, |
| bos_token="<|begin_of_text|>", |
| - eos_token="<|end_of_text|>", |
| - chat_template=chat_template, |
| + eos_token="<|end_of_text|>" if not instruct else "<|eot_id|>", |
| + chat_template=chat_template if instruct else None, |
| model_input_names=["input_ids", "attention_mask"], |
| + model_max_length=model_max_length, |
| ) |
| |
| |
| -def write_tokenizer(tokenizer_path, input_tokenizer_path, llama_version=2): |
| +def write_tokenizer(tokenizer_path, input_tokenizer_path, llama_version="2", special_tokens=None, instruct=False): |
| tokenizer_class = LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast |
| - if llama_version == 3: |
| - tokenizer = Llama3Converter(input_tokenizer_path).tokenizer |
| + if llama_version in ["3", "3.1"]: |
| + tokenizer = Llama3Converter( |
| + input_tokenizer_path, |
| + special_tokens, |
| + instruct, |
| + model_max_length=CONTEXT_LENGTH_FOR_VERSION[llama_version] |
| + ).tokenizer |
| else: |
| tokenizer = tokenizer_class(input_tokenizer_path) |
| print(f"Saving a {tokenizer_class.__name__} to {tokenizer_path}.") |
| @@ -362,6 +372,37 @@ def write_tokenizer(tokenizer_path, input_tokenizer_path, llama_version=2): |
| return tokenizer |
| |
| |
| +DEFAULT_LLAMA_SPECIAL_TOKENS = { |
| + "3": [ |
| + "<|begin_of_text|>", |
| + "<|end_of_text|>", |
| + "<|reserved_special_token_0|>", |
| + "<|reserved_special_token_1|>", |
| + "<|reserved_special_token_2|>", |
| + "<|reserved_special_token_3|>", |
| + "<|start_header_id|>", |
| + "<|end_header_id|>", |
| + "<|reserved_special_token_4|>", |
| + "<|eot_id|>", # end of turn |
| + ] |
| + + [f"<|reserved_special_token_{i}|>" for i in range(5, 256 - 5)], |
| + "3.1": [ |
| + "<|begin_of_text|>", |
| + "<|end_of_text|>", |
| + "<|reserved_special_token_0|>", |
| + "<|reserved_special_token_1|>", |
| + "<|finetune_right_pad_id|>", |
| + "<|reserved_special_token_2|>", |
| + "<|start_header_id|>", |
| + "<|end_header_id|>", |
| + "<|eom_id|>", # end of message |
| + "<|eot_id|>", # end of turn |
| + "<|python_tag|>", |
| + ] |
| + + [f"<|reserved_special_token_{i}|>" for i in range(3, 256 - 8)], |
| +} |
| + |
| + |
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument( |
| @@ -383,9 +424,9 @@ def main(): |
| # Different Llama versions used different default values for max_position_embeddings, hence the need to be able to specify which version is being used. |
| parser.add_argument( |
| "--llama_version", |
| - choices=[1, 2, 3], |
| - default=1, |
| - type=int, |
| + choices=["1", "2", "3", "3.1"], |
| + default="1", |
| + type=str, |
| help="Version of the Llama model to convert. Currently supports Llama1 and Llama2. Controls the context size", |
| ) |
| parser.add_argument( |
| @@ -394,11 +435,34 @@ def main(): |
| type=int, |
| help="The number of individual shards used for the model. Does not have to be the same as the number of consolidated_xx.pth", |
| ) |
| + parser.add_argument( |
| + "--special_tokens", |
| + default=None, |
| + type=List[str], |
| + help="The list of special tokens that should be added to the model.", |
| + ) |
| + parser.add_argument( |
| + "--instruct", |
| + default=False, |
| + type=bool, |
| + help="Whether the model is an instruct model or not. Will affect special tokens for llama 3.1.", |
| + ) |
| args = parser.parse_args() |
| if args.model_size is None and args.num_shards is None: |
| raise ValueError("You have to set at least `num_shards` if you are not giving the `model_size`") |
| + if args.special_tokens is None: |
| + args.special_tokens = DEFAULT_LLAMA_SPECIAL_TOKENS[str(args.llama_version)] |
| + |
| spm_path = os.path.join(args.input_dir, "tokenizer.model") |
| - vocab_size = len(write_tokenizer(args.output_dir, spm_path, llama_version=args.llama_version)) |
| + vocab_size = len( |
| + write_tokenizer( |
| + args.output_dir, |
| + spm_path, |
| + llama_version=args.llama_version, |
| + special_tokens=args.special_tokens, |
| + instruct=args.instruct |
| + ) |
| + ) |
| if args.model_size != "tokenizer_only": |
| write_model( |
| model_path=args.output_dir, |
| @@ -408,6 +472,7 @@ def main(): |
| llama_version=args.llama_version, |
| vocab_size=vocab_size, |
| num_shards=args.num_shards, |
| + instruct=args.instruct |
| ) |
| |
| |
|
|
