Fix max_position_embeddings default value for llama2 to 4096 #28241

src/transformers/models/llama/configuration_llama.py

@@ -60,7 +60,7 @@ class LlamaConfig(PretrainedConfig):
             `num_attention_heads`.
         hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
             The non-linear activation function (function or string) in the decoder.
-        max_position_embeddings (`int`, *optional*, defaults to 2048):
+        max_position_embeddings (`int`, *optional*, defaults to 4096):
             The maximum sequence length that this model might ever be used with. Llama 1 supports up to 2048 tokens,
             Llama 2 up to 4096, CodeLlama up to 16384.
         initializer_range (`float`, *optional*, defaults to 0.02):

src/transformers/models/llama/convert_llama_weights_to_hf.py

@@ -80,7 +80,9 @@ def write_json(text, path):
         json.dump(text, f)
 
 
-def write_model(model_path, input_base_path, model_size, tokenizer_path=None, safe_serialization=True):
+def write_model(
+    model_path, input_base_path, model_size, tokenizer_path=None, safe_serialization=True, llama_version=1
+):
     # for backward compatibility, before you needed the repo to be called `my_repo/model_size`
     if not os.path.isfile(os.path.join(input_base_path, "params.json")):
         input_base_path = os.path.join(input_base_path, model_size)
@@ -102,7 +104,16 @@ def write_model(model_path, input_base_path, model_size, tokenizer_path=None, sa
     if base > 10000.0:
         max_position_embeddings = 16384
     else:
-        max_position_embeddings = 2048
+        # 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
+        else:
+            raise NotImplementedError(
+                f"Version {llama_version} of llama is not supported yet. "
+                "Current supported versions of llama are [1, 2]."
+            )
 
     tokenizer_class = LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
     if tokenizer_path is not None:
@@ -301,6 +312,14 @@ def main():
         help="Location to write HF model and tokenizer",
     )
     parser.add_argument("--safe_serialization", type=bool, help="Whether or not to save using `safetensors`.")
+    # 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],
+        default=1,
+        type=int,
+        help="Version of the Llama model to use. Current supported versions are Llama1 and Llama2.",
+    )
     args = parser.parse_args()
     spm_path = os.path.join(args.input_dir, "tokenizer.model")
     if args.model_size != "tokenizer_only":
@@ -310,6 +329,7 @@ def main():
             model_size=args.model_size,
             safe_serialization=args.safe_serialization,
             tokenizer_path=spm_path,
+            llama_version=args.llama_version,
         )
     else:
         write_tokenizer(args.output_dir, spm_path)