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)